perf: ⚡️ change find_spec -> classmethod

ZhaoQi99 · Nov 5, 2024 · b1b51b9 · b1b51b9
1 parent 32db20c
commit b1b51b9
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/pyencrypt/encrypt.py b/pyencrypt/encrypt.py
@@ -33,8 +33,6 @@ def _encrypt_file(
 def can_encrypt(path: Path) -> bool:
     if path.name in NOT_ALLOWED_ENCRYPT_FILES:
         return False
-    if "management/commands/" in path.as_posix():
-        return False
     if path.suffix != ".py":
         return False
     return True

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",
@@ -73,17 +76,23 @@ def get_data(self, path: _Path) -> bytes:
 
 
 class EncryptFileFinder(abc.MetaPathFinder, Base):
+    @classmethod
     def find_spec(
-        self, fullname: str, path: Sequence[_Path], target: types.ModuleType = None
+        cls, fullname: str, path: Sequence[_Path], target: types.ModuleType = None
     ) -> 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()
@@ -92,6 +101,11 @@ def find_spec(
         loader = EncryptFileLoader(file_path)
         return spec_from_loader(name=fullname, loader=loader, origin="origin-encrypt")
 
+    @classmethod
+    def invalidate_caches(cls):
+        pass
+
 
 # TODO: generate randomly AES Class
-sys.meta_path.insert(0, EncryptFileFinder())
+machinery.EXTENSION_SUFFIXES.append(ENCRYPT_SUFFIX)
+sys.meta_path.insert(0, EncryptFileFinder)