From 57272d337e04dc343057246ec614836774b2a1ed Mon Sep 17 00:00:00 2001
From: Qi Zhao <zhaoqi99@outlook.com>
Date: Fri, 29 Nov 2024 15:43:01 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20SyntaxError=20when=20call?=
 =?UTF-8?q?=20`traceback`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

✅ Closes: #49
---
 pyencrypt/loader.py     | 13 +++++++
 tests/conftest.py       |  4 +--
 tests/test_traceback.py | 76 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_traceback.py

diff --git a/pyencrypt/loader.py b/pyencrypt/loader.py
index a4f6bc2..0d87a2c 100644
--- a/pyencrypt/loader.py
+++ b/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")
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 260ab0e..34d2a04 100644
--- a/tests/conftest.py
+++ b/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)
 
 
diff --git a/tests/test_traceback.py b/tests/test_traceback.py
new file mode 100644
index 0000000..9d807bb
--- /dev/null
+++ b/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 as e:
+        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 as e:
+        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