improvements to defer

thirdparty/defer/defer.py

@@ -2,7 +2,6 @@ import sys
 from typing import Any, Callable, ParamSpec, TypeVar
 
 from defer._defer import _Defer
-from defer.sugar._parse import _ParseDefer
 
 P = ParamSpec("P")
 T = TypeVar("T")
@@ -11,7 +10,7 @@ T = TypeVar("T")
 class Defer:
     def __call__(self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs):
         if sys.gettrace() is None:
-            sys.settrace(_ParseDefer.IDENTITY)
+            sys.settrace(_DeferTrace())
         frame = sys._getframe(1)
         if not isinstance(frame.f_trace, _Defer):
             frame.f_trace = _Defer(frame.f_trace)
@@ -21,3 +20,9 @@ class Defer:
     def __contains__(self, fn: Any):
         breakpoint()
         self(fn)
+
+
+class _DeferTrace:
+    """Minimal no-op trace for when @defer is used but no other trace is active."""
+    def __call__(self, frame, event, arg):
+        return self

thirdparty/defer/sugar/__init__.py

@@ -1,12 +1,12 @@
 import sys
 
 from defer.defer import Defer
-from defer.sugar._parse import _ParseDefer
 
 
 def install():
-    if not isinstance(sys.gettrace(), _ParseDefer):
+    """Only loads markers - does NOT install trace. Trace is lazy."""
-        sys.settrace(_ParseDefer(sys.gettrace()))
+    from defer.sugar._parse import _load_defer_markers
+    _load_defer_markers()
 
 
 defer = Defer()

thirdparty/defer/sugar/_parse.py

@@ -1,108 +1,56 @@
 import sys
 import os
+import ast
 from ast import (
     AsyncFunctionDef,
     FunctionDef,
 )
-from collections import deque
 from types import FrameType
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, Optional
 
 from defer.errors import FreeVarsError
 from defer.sugar.transformer import RewriteDefer
 
 
-def _is_dunder(name: str) -> bool:
+_DEFER_MARKER = "# defer: parse"
-    return name.startswith("__") and name.endswith("__")
+_enabled_files: set[str] = set()
 
 
-_SITE_PACKAGES = "site-packages"
+def _load_defer_markers():
-_VENV = ".venv"
+    global _enabled_files
-_THIRDPARTY = "thirdparty"
+    _enabled_files.clear()
-_SRC_WIN = "\\src\\"
-_SRC_UNIX = "/src/"
-
-
-def _is_our_code(filename: str) -> bool:
-    if not filename:
-        return False
-    if _SRC_WIN in filename or _SRC_UNIX in filename:
-        return True
-    return False
-
-
-class _ParseDefer:
-    IDENTITY = lambda *_: None  # noqa: E731
-
-    def __init__(self, tracefn: Optional[Callable]) -> None:
-        self.tracefn = tracefn or self.IDENTITY
-        self.pending: deque[Executing] = deque()
-
-    def __call__(self, frame: FrameType, event: str, arg: Any):
-        self.tracefn(frame, event, arg)
-
     try:
-            filename = frame.f_code.co_filename
+        for root, dirs, filenames in os.walk('src'):
-            if not filename:
+            for filename in filenames:
-                return self
+                if filename.endswith('.py'):
-            if _SITE_PACKAGES in filename or _VENV in filename or _THIRDPARTY in filename:
+                    relpath = os.path.join(root, filename)
-                return self
+                    abspath = os.path.abspath(relpath)
-            if not _is_our_code(filename):
+                    try:
-                return self
+                        with open(abspath, "r", encoding="utf-8", errors="ignore") as f:
+                            for i in range(10):
+                                line = f.readline()
+                                if not line:
+                                    break
+                                if _DEFER_MARKER in line:
+                                    _enabled_files.add(abspath)
+                                    break
+                    except Exception:
+                        pass
     except Exception:
         pass
 
-        if event != "line":
-            return self
 
-        try:
+def install():
-            from executing.executing import Executing, Source
+    """Only loads markers - does NOT install trace. Trace is lazy."""
-        except ImportError:
+    _load_defer_markers()
-            return self
 
-        try:
-            exc = Source.executing(frame)
-        except Exception:
-            return self
 
-        if not (stmt := next(iter(exc.statements), None)):
+class _ParseDefer:
-            return self
+    """Lazy defer - trace only activates when @defer is actually used."""
-        if isinstance(stmt, (AsyncFunctionDef, FunctionDef)):
+    IDENTITY = lambda *_: None
-            if _is_dunder(stmt.name):
-                return self
-            self.pending.append(exc)
-            return self
-        if not self.pending or frame.f_back is not self.pending[-1].frame.f_back:
-            return self
 
-        stmts = self.pending.pop().statements
+    def __init__(self, tracefn: Optional[Callable]) -> None:
-        node = cast(FunctionDef | AsyncFunctionDef, next(iter(stmts)))
+        self.tracefn = tracefn or self.IDENTITY
-        fn_name = node.name
 
-        if fn_name not in frame.f_locals:
+    def __call__(self, frame: FrameType, event: str, arg: Any):
-            return self
+        return self.tracefn(frame, event, arg)
-        fn = frame.f_locals[fn_name]
-
-        if not callable(fn):
-            return self
-        try:
-            code = fn.__code__
-        except AttributeError:
-            return self
-
-        try:
-            mod = fn.__module__
-        except AttributeError:
-            return self
-        if mod in sys.stdlib_module_names:
-            return self
-        if code.co_freevars:
-            return self
-        if not (ast := RewriteDefer.transform(node)):
-            return self
-
-        locals = frame.f_locals.copy()
-        del locals[fn_name]
-        exec(compile(ast, frame.f_code.co_filename, "exec"), frame.f_globals, locals)
-        frame.f_locals[fn_name].__code__ = locals[fn_name].__code__
-        return self