LazyModule for lazily proxying module imports.

common/lib/codejail/codejail/lazymod.py

@@ -0,0 +1,41 @@
+"""A module proxy for delayed importing of modules.
+
+Lifted from http://barnesc.blogspot.com/2006/06/automatic-python-imports-with-autoimp.html
+
+"""
+
+import sys
+
+class LazyModule(object):
+    """A lazy module proxy."""
+
+    def __init__(self, modname):
+        self.__dict__['__name__'] = modname
+        self._set_mod(None)
+
+    def _set_mod(self, mod):
+        if mod is not None:
+            self.__dict__ = mod.__dict__
+        self.__dict__['_lazymod_mod'] = mod
+
+    def _load_mod(self):
+        __import__(self.__name__)
+        self._set_mod(sys.modules[self.__name__])
+
+    def __getattr__(self, name):
+        if self.__dict__['_lazymod_mod'] is None:
+            self._load_mod()
+
+        mod = self.__dict__['_lazymod_mod']
+
+        if hasattr(mod, name):
+            return getattr(mod, name)
+        else:
+            try:
+                subname = '%s.%s' % (self.__name__, name)
+                __import__(subname)
+                submod = getattr(mod, name)
+            except ImportError:
+                raise AttributeError("'module' object has no attribute %r" % name)
+            self.__dict__[name] = LazyModule(subname, submod)
+            return self.__dict__[name]

common/lib/codejail/codejail/tests/test_lazymod.py

@@ -0,0 +1,26 @@
+"""Test lazymod.py"""
+
+import sys
+import unittest
+
+from codejail.lazymod import LazyModule
+from codejail.util import ModuleIsolation
+
+
+class TestLazyMod(unittest.TestCase):
+
+    def setUp(self):
+        # Each test will remove modules that it imported.
+        self.addCleanup(ModuleIsolation().clean_up)
+
+    def test_simple(self):
+        # Import some stdlib module that has not been imported before
+        self.assertNotIn("colorsys", sys.modules)
+        colorsys = LazyModule("colorsys")
+        hsv = colorsys.rgb_to_hsv(.3, .4, .2)
+        self.assertEqual(hsv[0], 0.25)
+
+    def test_dotted(self):
+        self.assertNotIn("email.utils", sys.modules)
+        email_utils = LazyModule("email.utils")
+        self.assertEqual(email_utils.quote('"hi"'), r'\"hi\"')

common/lib/codejail/codejail/util.py

@@ -3,8 +3,10 @@
 import contextlib
 import os
 import shutil
+import sys
 import tempfile
 
+
 class TempDirectory(object):
     def __init__(self, delete_when_done=True):
         self.delete_when_done = delete_when_done
@@ -28,3 +30,22 @@ def temp_directory(delete_when_done=True):
         yield tmp.temp_dir
     finally:
         tmp.clean_up()
+
+
+class ModuleIsolation(object):
+    """
+    Manage changes to sys.modules so that we can roll back imported modules.
+
+    Create this object, it will snapshot the currently imported modules. When
+    you call `clean_up()`, it will delete any module imported since its creation.
+    """
+    def __init__(self):
+        # Save all the names of all the imported modules.
+        self.mods = set(sys.modules)
+
+    def clean_up(self):
+        # Get a list of modules that didn't exist when we were created
+        new_mods = [m for m in sys.modules if m not in self.mods]
+        # and delete them all so another import will run code for real again.
+        for m in new_mods:
+            del sys.modules[m]