From 6af1a64cce0d35ff153459420010d3987cdb9d44 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 20 Jan 2019 15:53:11 +0000
Subject: [PATCH] master: handle crazy non-modules in sys.modules again; closes
 #310.

---
 docs/changelog.rst                      |  6 +++++
 mitogen/master.py                       | 22 ++++++++++++++++-
 tests/data/pkg_like_plumbum/__init__.py |  0
 tests/data/pkg_like_plumbum/colors.py   | 11 +++++++++
 tests/module_finder_test.py             | 33 +++++++++++++++++++++++++
 5 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 tests/data/pkg_like_plumbum/__init__.py
 create mode 100644 tests/data/pkg_like_plumbum/colors.py

diff --git a/docs/changelog.rst b/docs/changelog.rst
index b3edb34e..5ef62611 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -246,6 +246,11 @@ Core Library
   causing Mitogen to install itself at the end of the importer chain rather
   than the front.
 
+* `#310 <https://github.com/dw/mitogen/issues/310>`_: support has returned for
+  trying to figure out the real source of non-module objects installed in
+  :data:`sys.modules`, so they can be imported. This is needed to handle syntax
+  sugar used by packages like :mod:`plumbum`.
+
 * `#349 <https://github.com/dw/mitogen/issues/349>`_: an incorrect format
   string could cause large stack traces when attempting to import built-in
   modules on Python 3.
@@ -387,6 +392,7 @@ bug reports, testing, features and fixes in this release contributed by
 `Duane Zamrok <https://github.com/dewthefifth>`_,
 `Eric Chang <https://github.com/changchichung>`_,
 `Guy Knights <https://github.com/knightsg>`_,
+`Jesse London <https://github.com/jesteria>`_,
 `Jiří Vávra <https://github.com/Houbovo>`_,
 `Jonathan Rosser <https://github.com/jrosser>`_,
 `Johan Beisser <https://github.com/jbeisser>`_,
diff --git a/mitogen/master.py b/mitogen/master.py
index 28d3fa9d..a1ed2dfb 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -453,8 +453,28 @@ class ModuleFinder(object):
 
         return path, source, is_pkg
 
+    def _get_module_via_parent_enumeration(self, fullname):
+        """
+        Attempt to fetch source code by examining the module's (hopefully less
+        insane) parent package. Required for older versions of
+        ansible.compat.six and plumbum.colors.
+        """
+        pkgname, _, modname = fullname.rpartition('.')
+        pkg = sys.modules.get(pkgname)
+        if pkg is None or not hasattr(pkg, '__file__'):
+            return
+
+        pkg_path = os.path.dirname(pkg.__file__)
+        try:
+            fp, path, ext = imp.find_module(modname, [pkg_path])
+            return path, fp.read().encode('utf-8'), False
+        except ImportError:
+            e = sys.exc_info()[1]
+            LOG.debug('imp.find_module(%r, %r) -> %s', modname, [pkg_path], e)
+
     get_module_methods = [_get_module_via_pkgutil,
-                          _get_module_via_sys_modules]
+                          _get_module_via_sys_modules,
+                          _get_module_via_parent_enumeration]
 
     def get_module_source(self, fullname):
         """Given the name of a loaded module `fullname`, attempt to find its
diff --git a/tests/data/pkg_like_plumbum/__init__.py b/tests/data/pkg_like_plumbum/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/data/pkg_like_plumbum/colors.py b/tests/data/pkg_like_plumbum/colors.py
new file mode 100644
index 00000000..b0fc4ce6
--- /dev/null
+++ b/tests/data/pkg_like_plumbum/colors.py
@@ -0,0 +1,11 @@
+
+import sys
+
+
+class EvilObject(object):
+    """
+    Wild cackles! I have come to confuse perplex your importer with rainbows!
+    """
+
+sys.modules[__name__] = EvilObject()
+
diff --git a/tests/module_finder_test.py b/tests/module_finder_test.py
index 1d5a0796..efc60af0 100644
--- a/tests/module_finder_test.py
+++ b/tests/module_finder_test.py
@@ -105,6 +105,39 @@ class GetModuleViaSysModulesTest(testlib.TestCase):
         self.assertIsNone(tup)
 
 
+class GetModuleViaParentEnumerationTest(testlib.TestCase):
+    klass = mitogen.master.ModuleFinder
+
+    def call(self, fullname):
+        return self.klass()._get_module_via_parent_enumeration(fullname)
+
+    def test_main_fails(self):
+        import __main__
+        self.assertIsNone(self.call('__main__'))
+
+    def test_dylib_fails(self):
+        # _socket comes from a .so
+        import _socket
+        tup = self.call('_socket')
+        self.assertIsNone(tup)
+
+    def test_builtin_fails(self):
+        # sys is built-in
+        tup = self.call('sys')
+        self.assertIsNone(tup)
+
+    def test_plumbum_colors_like_pkg_succeeds(self):
+        # plumbum has been eating too many rainbow-colored pills
+        import pkg_like_plumbum.colors
+        path, src, is_pkg = self.call('pkg_like_plumbum.colors')
+        self.assertEquals(path,
+            testlib.data_path('pkg_like_plumbum/colors.py'))
+
+        s = open(testlib.data_path('pkg_like_plumbum/colors.py')).read()
+        self.assertEquals(mitogen.core.to_text(src), s)
+        self.assertFalse(is_pkg)
+
+
 class ResolveRelPathTest(testlib.TestCase):
     klass = mitogen.master.ModuleFinder