fix: 🐛 add ENCRYPT_SUFFIX for inspect.getmodulename

ZhaoQi99 · Nov 5, 2024 · a071ad1 · a071ad1
1 parent 7e08edb
commit a071ad1
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/pyencrypt/loader.py b/pyencrypt/loader.py
@@ -2,7 +2,7 @@
 import sys
 import traceback
 import types
-from importlib import abc
+from importlib import abc, machinery
 from importlib._bootstrap_external import _NamespacePath
 from importlib.machinery import ModuleSpec
 from importlib.util import spec_from_loader
@@ -21,6 +21,9 @@ def __dir__(self) -> Iterable[str]:
         return []
 
 
+ENCRYPT_SUFFIX = ".pye"
+
+
 class EncryptFileLoader(abc.SourceLoader, Base):
     POSSIBLE_PATH = [
         Path(os.path.expanduser("~")) / ".licenses" / "license.lic",
@@ -79,12 +82,17 @@ def find_spec(
     ) -> ModuleSpec:
         if path:
             if isinstance(path, _NamespacePath):
-                file_path = Path(path._path[0]) / f'{fullname.rsplit(".", 1)[-1]}.pye'
+                file_path = (
+                    Path(path._path[0])
+                    / f'{fullname.rsplit(".", 1)[-1]}{ENCRYPT_SUFFIX}'
+                )
             else:
-                file_path = Path(path[0]) / f'{fullname.rsplit(".", 1)[-1]}.pye'
+                file_path = (
+                    Path(path[0]) / f'{fullname.rsplit(".", 1)[-1]}{ENCRYPT_SUFFIX}'
+                )
         else:
             for p in sys.path:
-                file_path = Path(p) / f"{fullname}.pye"
+                file_path = Path(p) / f"{fullname}{ENCRYPT_SUFFIX}"
                 if file_path.exists():
                     break
         file_path = file_path.absolute().as_posix()
@@ -99,4 +107,5 @@ def invalidate_caches(cls):
 
 
 # TODO: generate randomly AES Class
+machinery.EXTENSION_SUFFIXES.append(ENCRYPT_SUFFIX)
 sys.meta_path.insert(0, EncryptFileFinder)