issue #590: teach importer to handle self-replacing modules

docs/changelog.rst

@@ -30,6 +30,12 @@ Enhancements
   <https://docs.ansible.com/ansible/latest/plugins/become.html>`_
   functionality, which will be addressed in a future release.
 
+Fixes
+^^^^^
+
+* `#590 <https://github.com/dw/mitogen/issues/590>`_: the importer can handle
+  modules that replace themselves in :mod:`sys.modules` during import.
+
 
 Thanks!
 ~~~~~~~

mitogen/core.py

@@ -1355,7 +1355,10 @@ class Importer(object):
             exec(code, vars(mod))
         else:
             exec('exec code in vars(mod)')
-        return mod
+
+        # #590: if a module replaces itself in sys.modules during import, below
+        # is necessary. This matches PyImport_ExecCodeModuleEx()
+        return sys.modules.get(fullname, mod)
 
     def get_filename(self, fullname):
         if fullname in self._cache:

tests/data/simple_pkg/imports_replaces_self.py

@@ -0,0 +1,6 @@
+# issue #590: this module imports a module that replaces itself in sys.modules
+# during initialization.
+import simple_pkg.replaces_self
+
+def subtract_one(n):
+    return simple_pkg.replaces_self.subtract_one(n)

tests/data/simple_pkg/replaces_self.py

@@ -0,0 +1,4 @@
+# issue #590: this module replaces itself in sys.modules during initialization.
+import sys
+import simple_pkg.b
+sys.modules[__name__] = simple_pkg.b

tests/importer_test.py

@@ -12,6 +12,7 @@ import mitogen.utils
 from mitogen.core import b
 
 import testlib
+import simple_pkg.imports_replaces_self
 
 
 class ImporterMixin(testlib.RouterMixin):
@@ -214,5 +215,13 @@ class Python24LineCacheTest(testlib.TestCase):
     pass
 
 
+class SelfReplacingModuleTest(testlib.RouterMixin, testlib.TestCase):
+    # issue #590
+    def test_importer_handles_self_replacement(self):
+        c = self.router.local()
+        self.assertEquals(0,
+            c.call(simple_pkg.imports_replaces_self.subtract_one, 1))
+
+
 if __name__ == '__main__':
     unittest2.main()