fix: 🐛 SyntaxError when call traceback

✅ Closes: #49

Author: ZhaoQi99

pyencrypt/loader.py

@@ -1,3 +1,4 @@
+import linecache
 import os
 import sys
 import traceback
@@ -64,6 +65,9 @@ def check(self) -> bool:
     def get_filename(self, fullname: str) -> str:
         return self.path
 
+    def get_source(self, fullname: str):
+        return None
+
     def get_data(self, path: _Path) -> bytes:
         try:
             __n, __d = self.__private_key.split("O", 1)
@@ -76,6 +80,12 @@ def get_data(self, path: _Path) -> bytes:
 
 
 class EncryptFileFinder(abc.MetaPathFinder, Base):
+    @staticmethod
+    def _cache_line(file_path):
+        stat = os.stat(file_path)
+        size, mtime = stat.st_size, stat.st_mtime
+        linecache.cache[file_path] = (size, mtime, [], file_path)
+
     @classmethod
     def find_spec(
         cls, fullname: str, path: Sequence[_Path], target: types.ModuleType = None
@@ -98,6 +108,9 @@ def find_spec(
         file_path = file_path.absolute().as_posix()
         if not os.path.exists(file_path):
             return None
+
+        cls._cache_line(file_path)
+
         loader = EncryptFileLoader(file_path)
         return spec_from_loader(name=fullname, loader=loader, origin="origin-encrypt")
 

tests/conftest.py

@@ -1,4 +1,5 @@
 import pytest
+
 from pyencrypt.encrypt import encrypt_file, encrypt_key, generate_so_file
 from pyencrypt.generate import generate_aes_key
 from pyencrypt.license import generate_license_file
@@ -47,7 +48,7 @@ def {function_name}():
     # generate loader.so
     key = generate_aes_key()
     new_path = file_path.with_suffix(".pye")
-    encrypt_file(file_path, key, new_path=new_path)
+    encrypt_file(file_path, key.decode(), new_path=new_path)
     file_path.unlink()
     cipher_key, d, n = encrypt_key(key)
     loader_path = generate_so_file(cipher_key, d, n, file_path.parent, license=license)
@@ -58,7 +59,6 @@ def {function_name}():
 
     # License
     license and generate_license_file(key.decode(), work_dir, **kwargs)
-    generate_license_file(key.decode(), work_dir)
     return (new_path, loader_path)
 
 

tests/test_traceback.py

@@ -0,0 +1,76 @@
+import sys
+from pathlib import Path
+from typing import Tuple
+
+import pytest
+
+CODE1 = """
+    try:
+        a = 1/0
+    except Exception as e:
+        import traceback
+        return traceback.format_exc()
+def fun():
+    return _fun()
+"""
+
+
+@pytest.mark.file(name="file_trace1", function="_fun", code=CODE1)
+def test_traceback_format_exc(file_and_loader: Tuple[Path], monkeypatch):
+    file_path, loader_path = file_and_loader
+    monkeypatch.syspath_prepend(loader_path.parent.as_posix())
+    monkeypatch.syspath_prepend(file_path.parent.as_posix())
+
+    sys.modules.pop("loader", None)
+
+    import loader
+    from file_trace1 import fun
+
+    msg = None
+    try:
+        msg = fun()
+    except Exception:
+        pass
+    if not msg:
+        pytest.fail("`format_exc` return None")
+
+
+CODE2 = """
+    try:
+        a = 1/0
+    except Exception as e:
+        import traceback
+        traceback.print_stack()
+def _fun2():
+    _fun()
+def _fun3():
+    _fun2()
+
+def fun2():
+    _fun3()
+"""
+
+
+@pytest.mark.file(name="file_trace2", function="_fun", code=CODE2)
+def test_traceback_print_stack(file_and_loader: Tuple[Path], monkeypatch, capsys):
+    file_path, loader_path = file_and_loader
+    monkeypatch.syspath_prepend(loader_path.parent.as_posix())
+    monkeypatch.syspath_prepend(file_path.parent.as_posix())
+
+    sys.modules.pop("loader", None)
+
+    import loader
+    from file_trace2 import fun2
+
+    try:
+        fun2()
+    except Exception:
+        pass
+
+    captured = capsys.readouterr()
+    out = captured.err
+
+    assert "in fun" in out
+    assert "in _fun3" in out
+    assert "in _fun2" in out
+    assert "in _fun" in out