Merge pull request #108 from dx-77/issue-107

fix: Change PYE finder (hook) logic from meta hook to path hook

Author: Falldog

src/pyconcrete/__init__.py

@@ -16,52 +16,32 @@
 
 import imp
 import marshal
+import struct
 import sys
-from os.path import exists, isdir, join
+from importlib._bootstrap_external import _get_supported_file_loaders
+from importlib.machinery import FileFinder, SourceFileLoader, SOURCE_SUFFIXES
+from . import _pyconcrete  # noqa: E402
+
 
-EXT_PY = '.py'
-EXT_PYC = '.pyc'
-EXT_PYD = '.pyd'
 EXT_PYE = '.pye'
 
 __all__ = ["info"]
 
-from . import _pyconcrete  # noqa: E402
 
 info = _pyconcrete.info
 encrypt_file = _pyconcrete.encrypt_file
 decrypt_file = _pyconcrete.decrypt_file
 decrypt_buffer = _pyconcrete.decrypt_buffer
 
+# We need to modify SOURCE_SUFFIXES, because it used in importlib.machinery.all_suffixes function which
+# called by inspect.getmodulename and we need to be able to detect the module name relative to .pye files
+# because .py can be deleted by us
+SOURCE_SUFFIXES.append(EXT_PYE)
 
-class PyeLoader(object):
-    def __init__(self, is_pkg, pkg_path, full_path):
-        self.is_pkg = is_pkg
-        self.pkg_path = pkg_path
-        self.full_path = full_path
-        with open(full_path, 'rb') as f:
-            self.data = f.read()
-
-    def new_module(self, fullname, path, package_path):
-        m = imp.new_module(fullname)
-        m.__file__ = path
-        m.__loader__ = self
-        if self.is_pkg:
-            m.__path__ = [package_path]
-
-        if "__name__" not in m.__dict__:
-            m.__name__ = fullname
-
-        return m
-
-    def load_module(self, fullname):
-        if fullname in sys.modules:  # skip reload by now ...
-            return sys.modules[fullname]
-
-        data = decrypt_buffer(self.data)  # decrypt pye
-
-        self._validate_version(data)
 
+class PyeLoader(SourceFileLoader):
+    @property
+    def magic(self):
         if sys.version_info >= (3, 7):
             # reference python source code
             # python/Lib/importlib/_bootstrap_external.py _code_to_timestamp_pyc() & _code_to_hash_pyc()
@@ -77,47 +57,39 @@ def load_module(self, fullname):
             # reference http://stackoverflow.com/questions/1830727/how-to-load-compiled-python-modules-from-memory
             # MAGIC + TIMESTAMP
             magic = 8
-
-        code = marshal.loads(data[magic:])
-
-        m = self.new_module(fullname, self.full_path, self.pkg_path)
-        sys.modules[fullname] = m
-        exec(code, m.__dict__)
-        return m
-
-    def is_package(self, fullname):
-        return self.is_pkg
+        return magic
 
     @staticmethod
     def _validate_version(data):
         magic = imp.get_magic()
         ml = len(magic)
         if data[:ml] != magic:
-            import struct
-
             # convert little-endian byte string to unsigned short
             py_magic = struct.unpack('<H', magic[:2])[0]
             pye_magic = struct.unpack('<H', data[:2])[0]
             raise ValueError("Python version doesn't match with magic: python(%d) != pye(%d)" % (py_magic, pye_magic))
 
+    def get_code(self, fullname):
+        if not self.path.endswith(EXT_PYE):
+            return super().get_code(fullname)
+
+        path = self.get_filename(fullname)
+        data = decrypt_buffer(self.get_data(path))
+        self._validate_version(data)
+        return marshal.loads(data[self.magic:])
+
+    def get_source(self, fullname):
+        if self.path.endswith(EXT_PYE):
+            return None
+        return super().get_source(fullname)
+
 
-class PyeMetaPathFinder(object):
-    def find_module(self, fullname, path=None):
-        mod_name = fullname.split('.')[-1]
-        paths = path if path else sys.path
+loader_details = [(PyeLoader, SOURCE_SUFFIXES)] + _get_supported_file_loaders()
 
-        for trypath in paths:
-            mod_path = join(trypath, mod_name)
-            is_pkg = isdir(mod_path)
-            if is_pkg:
-                full_path = join(mod_path, '__init__' + EXT_PYE)
-                pkg_path = mod_path
-            else:
-                full_path = mod_path + EXT_PYE
-                pkg_path = trypath
 
-            if exists(full_path):
-                return PyeLoader(is_pkg, pkg_path, full_path)
+def install():
+    sys.path_importer_cache.clear()
+    sys.path_hooks.insert(0, FileFinder.path_hook(*loader_details))
 
 
-sys.meta_path.insert(0, PyeMetaPathFinder())
+install()