From 3435f24e8dc74db974ee1a33e2911ccaf7886f0c Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 27 Jan 2019 16:17:56 +0000
Subject: [PATCH] issue #479: ModuleFinder special case for __main__ on Py3.x.

---
 mitogen/master.py           | 41 ++++++++++++++++++++++++++++++++-----
 tests/module_finder_test.py | 27 ++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/mitogen/master.py b/mitogen/master.py
index aea28785..2930b42b 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -409,9 +409,37 @@ class ModuleFinder(object):
         if os.path.exists(path) and self._looks_like_script(path):
             return path
 
+    def _get_main_module_defective_python_3x(self, fullname):
+        """
+        Recent versions of Python 3.x introduced an incomplete notion of
+        importer specs, and in doing so created permanent asymmetry in the
+        :mod:`pkgutil` interface handling for the `__main__` module. Therefore
+        we must handle `__main__` specially.
+        """
+        if fullname != '__main__':
+            return None
+
+        mod = sys.modules.get(fullname)
+        if not mod:
+            return None
+
+        path = getattr(mod, '__file__', None)
+        if not (os.path.exists(path) and self._looks_like_script(path)):
+            return None
+
+        fp = open(path, 'rb')
+        try:
+            source = fp.read()
+        finally:
+            fp.close()
+
+        return path, source, False
+
     def _get_module_via_pkgutil(self, fullname):
-        """Attempt to fetch source code via pkgutil. In an ideal world, this
-        would be the only required implementation of get_module()."""
+        """
+        Attempt to fetch source code via pkgutil. In an ideal world, this would
+        be the only required implementation of get_module().
+        """
         try:
             # Pre-'import spec' this returned None, in Python3.6 it raises
             # ImportError.
@@ -543,9 +571,12 @@ class ModuleFinder(object):
         """
         self._found_cache[fullname] = (path, source, is_pkg)
 
-    get_module_methods = [_get_module_via_pkgutil,
-                          _get_module_via_sys_modules,
-                          _get_module_via_parent_enumeration]
+    get_module_methods = [
+        _get_main_module_defective_python_3x,
+        _get_module_via_pkgutil,
+        _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/module_finder_test.py b/tests/module_finder_test.py
index b6bf2111..bce3b70d 100644
--- a/tests/module_finder_test.py
+++ b/tests/module_finder_test.py
@@ -50,6 +50,33 @@ class IsStdlibNameTest(testlib.TestCase):
         self.assertFalse(self.func('mitogen.fakessh'))
 
 
+class GetMainModuleDefectivePython3x(testlib.TestCase):
+    klass = mitogen.master.ModuleFinder
+
+    def call(self, fullname):
+        return self.klass()._get_main_module_defective_python_3x(fullname)
+
+    def test_builtin(self):
+        self.assertEquals(None, self.call('sys'))
+
+    def test_not_main(self):
+        self.assertEquals(None, self.call('mitogen'))
+
+    def test_main(self):
+        import __main__
+
+        path, source, is_pkg = self.call('__main__')
+        self.assertTrue(path is not None)
+        self.assertTrue(os.path.exists(path))
+        self.assertEquals(path, __main__.__file__)
+        fp = open(path, 'rb')
+        try:
+            self.assertEquals(source, fp.read())
+        finally:
+            fp.close()
+        self.assertFalse(is_pkg)
+
+
 class GetModuleViaPkgutilTest(testlib.TestCase):
     klass = mitogen.master.ModuleFinder