New functions: uninstall(), is_installed()

polyloader/__init__.py

@@ -5,10 +5,9 @@
 __email__ = 'elf.sternberg@gmail.com'
 __version__ = '0.1.0'
 
-if sys.version_info[0:2] >= (2, 6):
-    from ._python2 import install, reset
-
 if sys.version_info[0] >= 3:
-    from ._python3 import install, reset
+    from ._python3 import install, uninstall, is_installed, reset
+elif sys.version_info[0:2] >= (2, 6):
+    from ._python2 import install, uninstall, is_installed, reset
 
-__all__ = ['install', 'reset']
+__all__ = ['install', 'uninstall', 'is_installed', 'reset']

polyloader/_python2.py

@@ -111,6 +111,21 @@ def _install(cls, compiler, suffixes):
             return
         cls._loader_handlers += [Loader(suf, compiler) for suf in suffixes]
 
+    @classmethod
+    def _uninstall(cls, suffixes):
+        if isinstance(suffixes, basestring):
+            suffixes = [suffixes]
+        suffixes = set(suffixes)
+        overlap = suffixes.intersection(set([suf[0] for suf in imp.get_suffixes()]))
+        if overlap:
+            raise RuntimeError("Removing a native Python extensions is not permitted.")
+
+        cls._loader_handlers = [ loader for loader in cls._loader_handlers if loader.suffix not in suffixes ]
+
+    @classmethod
+    def _is_installed(cls, suffix):
+        return any( loader.suffix == suffix for loader in cls._loader_handlers )
+
     @classmethod
     def getmodulename(cls, path):
         filename = os.path.basename(path)
@@ -174,6 +189,13 @@ def install(compiler, suffixes):
         PolyFinder._installed = True
     PolyFinder._install(compiler, suffixes)
 
+def uninstall(suffixes):
+    if not PolyFinder._installed:
+        return
+    PolyFinder._uninstall(suffixes)
+
+def is_installed(suffix):
+    return PolyFinder._is_installed(suffix)
 
 def reset():
     PolyFinder._loader_handlers = []

polyloader/_python3.py

@@ -152,6 +152,25 @@ def _install(cls, compiler, suffixes):
                              dict(_compiler = compiler))
         cls._custom_loaders += [(EXS + suffix, newloader) for suffix in suffixset]
 
+    @classmethod
+    def _uninstall(cls, suffixes):
+        if not suffixes:
+            return
+        if isinstance(suffixes, str):
+            suffixes = [suffixes]
+        suffixset = set(suffixes)
+        overlap = suffixset.intersection(set([suf[0] for suf in cls._native_loaders]))
+        if overlap:
+            raise RuntimeError("Removing a native Python extensions is not permitted.")
+
+        exs_suffixset = set( EXS + suffix for suffix in suffixset )
+        cls._custom_loaders = [ (suffix, loader) for suffix, loader in cls._custom_loaders if suffix not in exs_suffixset ]
+
+    @classmethod
+    def _is_installed(cls, suffix):
+        suffix = EXS + suffix
+        return any( suf == suffix for suf, loader in cls._custom_loaders )
+
     @classmethod
     def getmodulename(cls, path):
         filename = os.path.basename(path)
@@ -268,6 +287,11 @@ def install(compiler, suffixes):
 
     PolyFileFinder._install(compiler, suffixes)
 
+def uninstall(suffixes):
+    PolyFileFinder._uninstall(suffixes)
+
+def is_installed(suffix):
+    return PolyFileFinder._is_installed(suffix)
 
 def reset():
     PolyFileFinder._custom_loaders = []

tests_py2/test_polyloader.py

@@ -11,6 +11,7 @@
 import polyloader
 import copy
 import sys
+import pytest
 
 # Note that these compilers don't actually load much out of the
 # source files.  That's not the point.  The point is to show that the
@@ -28,6 +29,8 @@ def __enter__(self):
         self.meta_path = copy.copy(sys.meta_path)
         self.modules = copy.copy(sys.modules)
         self.path_importer_cache = copy.copy(sys.path_importer_cache)
+        self.dont_write_bytecode = sys.dont_write_bytecode
+        sys.dont_write_bytecode = True
         return sys
 
     def __exit__(self, type, value, traceback):
@@ -36,6 +39,7 @@ def __exit__(self, type, value, traceback):
         sys.meta_path = self.meta_path
         sys.modules = self.modules
         sys.path_importer_cache = self.path_importer_cache
+        sys.dont_write_bytecode = self.dont_write_bytecode
 
 
 class Compiler:
@@ -99,6 +103,52 @@ def test_import3(self):
             assert(result2 == "Success for 2: Test Two")
             assert(result3 == "Success for 3: Test Three")
 
+class Test_Polymorph_Reset(object):
+    def test_reset_after_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("2"), ['2'])
+            polyloader.install(compiler("3"), ['3'])
+            from tests_py2.polytestmix.test3 import result as result3
+            polyloader.reset()
+            with pytest.raises(ImportError):
+                from tests_py2.polytestmix.test2 import result as result2
+    def test_reset_before_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("3"), ['3'])
+            polyloader.reset()
+            with pytest.raises(ImportError):
+                from tests_py2.polytestmix.test3 import result as result3
+
+class Test_Polymorph_Uninstall(object):
+    def test_uninstall_after_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("2"), ['2'])
+            polyloader.install(compiler("3"), ['3'])
+            import tests_py2.polytestmix.test3
+            assert(polyloader.is_installed('2'))
+            assert(polyloader.is_installed('3'))
+            polyloader.uninstall('2')
+            assert(not polyloader.is_installed('2'))
+            assert(polyloader.is_installed('3'))
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test2
+            import tests_py2.polytestmix.test1
+    def test_uninstall_before_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("2"), ['2'])
+            polyloader.install(compiler("3"), ['3'])
+            assert(polyloader.is_installed('2'))
+            assert(polyloader.is_installed('3'))
+            polyloader.uninstall('2')
+            polyloader.uninstall('3')
+            assert(not polyloader.is_installed('2'))
+            assert(not polyloader.is_installed('3'))
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test2
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test3
+            import tests_py2.polytestmix.test1
+
 class Test_Polymorph_Iterator(object):
     ''' The Django Compatibility test: Can we load arbitrary modules from a package? '''
     def test_iterator(self):

tests_py3/test_polyloader.py

@@ -11,6 +11,7 @@
 import polyloader
 import copy
 import sys
+import pytest
 
 # Note that these compilers don't actually load much out of the
 # source files.  That's not the point.  The point is to show that the
@@ -28,6 +29,8 @@ def __enter__(self):
         self.meta_path = copy.copy(sys.meta_path)
         self.modules = copy.copy(sys.modules)
         self.path_importer_cache = copy.copy(sys.path_importer_cache)
+        self.dont_write_bytecode = sys.dont_write_bytecode
+        sys.dont_write_bytecode = True
         return sys
 
     def __exit__(self, type, value, traceback):
@@ -36,6 +39,7 @@ def __exit__(self, type, value, traceback):
         sys.meta_path = self.meta_path
         sys.modules = self.modules
         sys.path_importer_cache = self.path_importer_cache
+        sys.dont_write_bytecode = self.dont_write_bytecode
 
 
 class Compiler:
@@ -99,6 +103,44 @@ def test_import3(self):
             assert(result2 == "Success for 2: Test Two")
             assert(result3 == "Success for 3: Test Three")
 
+class Test_Polymorph_Reset(object):
+    def test_reset_after_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("2"), ['2'])
+            polyloader.install(compiler("3"), ['3'])
+            from tests_py2.polytestmix.test3 import result as result3
+            polyloader.reset()
+            with pytest.raises(ImportError):
+                from tests_py2.polytestmix.test2 import result as result2
+    def test_reset_before_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("3"), ['3'])
+            polyloader.reset()
+            with pytest.raises(ImportError):
+                from tests_py2.polytestmix.test3 import result as result3
+
+class Test_Polymorph_Uninstall(object):
+    def test_uninstall_after_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("3"), ['3'])
+            polyloader.install(compiler("2"), ['2'])
+            import tests_py2.polytestmix.test3
+            polyloader.uninstall('2')
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test2
+            import tests_py2.polytestmix.test1
+    def test_uninstall_before_import(self):
+        with ImportEnvironment() as sys:
+            polyloader.install(compiler("3"), ['3'])
+            polyloader.install(compiler("2"), ['2'])
+            polyloader.uninstall('3')
+            polyloader.uninstall('2')
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test3
+            with pytest.raises(ImportError):
+                import tests_py2.polytestmix.test2
+            import tests_py2.polytestmix.test1
+
 class Test_Polymorph_Iterator(object):
     ''' The Django Compatibility test: Can we load arbitrary modules from a package? '''
     def test_iterator(self):