fix?

Author: codeSamuraii

pyfuse/graph/analyzer.py

@@ -346,9 +346,15 @@ def _resolve_owner_class(qualname: str) -> str | None:
     if len(parts) == 1:
         return None
     prefix = parts[0]
-    if "<locals>" in prefix:
-        return None
-    return prefix
+    if "<locals>" not in prefix:
+        return prefix
+    # For nested classes like "outer.<locals>.MyClass.__init__",
+    # extract the class name after the last "<locals>." segment.
+    after_locals = prefix.rsplit("<locals>.", 1)[-1]
+    # If there's still a class name (not empty, not another scope marker)
+    if after_locals and "<" not in after_locals:
+        return after_locals
+    return None
 
 
 def _extract_annotation_type_names(annotation: ast.expr) -> set[str]:

pyfuse/graph/graph.py

@@ -128,28 +128,39 @@ def _try_get_lambda_source(func: Callable[..., object]) -> str | None:
 
 def _capture_closure(
     func: Callable[..., object],
-) -> tuple[dict[str, str], dict[str, str], dict[str, Callable[..., object]]]:
+) -> tuple[dict[str, str], dict[str, str], dict[str, Callable[..., object]], list[ImportInfo], dict[str, type]]:
     """Extract closure variables and traced function references from *func*.
 
-    Returns ``(closure_vars, closure_func_refs, closure_func_objects)``
+    Returns ``(closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes)``
     where *closure_vars* maps variable names to repr strings,
-    *closure_func_refs* maps variable names to qualified names, and
+    *closure_func_refs* maps variable names to qualified names,
     *closure_func_objects* maps qualified names to the actual callable
-    objects (for auto-registration of non-traced functions).
+    objects (for auto-registration of non-traced functions),
+    *closure_module_imports* is a list of :class:`ImportInfo` for module
+    objects found in the closure (from inline imports), and
+    *closure_classes* maps variable names to user-defined class objects
+    that need auto-registration.
     """
     closure_vars: dict[str, str] = {}
     closure_func_refs: dict[str, str] = {}
     closure_func_objects: dict[str, Callable[..., object]] = {}
+    closure_module_imports: list[ImportInfo] = []
+    closure_classes: dict[str, type] = {}
 
     if not func.__code__.co_freevars:
-        return closure_vars, closure_func_refs, closure_func_objects
+        return closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes
 
     try:
         closure_info = inspect.getclosurevars(func)
     except ValueError:
-        return closure_vars, closure_func_refs, closure_func_objects
+        return closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes
 
     for name, value in closure_info.nonlocals.items():
+        # Skip the implicit __class__ cell injected by Python for super() calls.
+        # Reconstructed code uses explicit super(ClassName, self) instead.
+        if name == "__class__" and inspect.isclass(value):
+            continue
+
         try:
             repr_value = repr(value)
         except Exception:
@@ -167,6 +178,17 @@ def _capture_closure(
         except SyntaxError:
             pass
 
+        # Module objects from inline imports (e.g. `import time as _time`)
+        if inspect.ismodule(value):
+            mod_name = value.__name__
+            if name == mod_name or name == mod_name.split(".")[0]:
+                stmt = f"import {mod_name}"
+            else:
+                stmt = f"import {mod_name} as {name}"
+            closure_module_imports.append(ImportInfo(statement=stmt, bound_name=name))
+            logger.debug("Closure var '%s' is module %s", name, mod_name)
+            continue
+
         # Callables: prefer source-level capture over serialization
         if getattr(value, "__pyfuse_traced__", False):
             unwrapped = value
@@ -189,6 +211,12 @@ def _capture_closure(
             logger.debug("Closure var '%s' is untraced user function %s", name, ref_qname)
             continue
 
+        # User-defined classes: auto-register all their methods
+        if inspect.isclass(value) and _is_user_class(value):
+            closure_classes[name] = value
+            logger.debug("Closure var '%s' is user class %s", name, value.__qualname__)
+            continue
+
         # Non-callable fallbacks
         ctor_expr = _try_constructor_expr(value)
         if ctor_expr is not None:
@@ -210,7 +238,7 @@ def _capture_closure(
             stacklevel=3,
         )
 
-    return closure_vars, closure_func_refs, closure_func_objects
+    return closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes
 
 
 def _mermaid_node_id(qname: str) -> str:
@@ -325,18 +353,27 @@ def register(self, func: Callable[..., object]) -> None:
                 "unavailable. Functions must be defined in .py source files."
             ) from exc
 
-        closure_vars, closure_func_refs, closure_func_objects = _capture_closure(original)
+        closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes = _capture_closure(original)
 
         for ref_qname, func_obj in closure_func_objects.items():
             if ref_qname not in self._nodes:
                 self._auto_register(func_obj)
 
+        for cls_obj in closure_classes.values():
+            self._auto_register_class(cls_obj)
+
+        # Add module imports from closures (inline imports like `import time as _time`)
+        existing = {imp.bound_name for imp in analysis.imports}
+        for imp in closure_module_imports:
+            if imp.bound_name not in existing:
+                analysis.imports.append(imp)
+                existing.add(imp.bound_name)
+
         if closure_vars:
             closure_names: set[str] = set()
             for cv in closure_vars.values():
                 closure_names |= get_used_names(cv)
             if closure_names:
-                existing = {imp.bound_name for imp in analysis.imports}
                 all_imports = get_module_imports(original)
                 for imp in all_imports:
                     if imp.bound_name in closure_names and imp.bound_name not in existing:
@@ -396,13 +433,40 @@ def _auto_register(self, func: Callable[..., object]) -> bool:
             )
             return False
 
+        closure_vars, closure_func_refs, closure_func_objects, closure_module_imports, closure_classes = _capture_closure(func)
+
+        # Add module imports from closures (inline imports)
+        existing_names = {imp.bound_name for imp in analysis.imports}
+        for imp in closure_module_imports:
+            if imp.bound_name not in existing_names:
+                analysis.imports.append(imp)
+                existing_names.add(imp.bound_name)
+
+        # Add imports needed by closure var expressions
+        if closure_vars:
+            closure_names: set[str] = set()
+            for cv in closure_vars.values():
+                closure_names |= get_used_names(cv)
+            if closure_names:
+                try:
+                    all_imports = get_module_imports(func)
+                    for imp in all_imports:
+                        if imp.bound_name in closure_names and imp.bound_name not in existing_names:
+                            analysis.imports.append(imp)
+                            existing_names.add(imp.bound_name)
+                except (OSError, TypeError):
+                    pass
+
         dependencies = [
             dep for dep in detect_traced_dependencies(
                 analysis.source, func.__module__, self._nodes,
                 owner_class=analysis.owner_class,
             )
             if dep != qualified_name
         ]
+        for ref_qname in closure_func_refs.values():
+            if ref_qname != qualified_name and ref_qname not in dependencies:
+                dependencies.append(ref_qname)
 
         node = FunctionNode(
             qualified_name=qualified_name,
@@ -412,12 +476,22 @@ def _auto_register(self, func: Callable[..., object]) -> bool:
             imports=analysis.imports,
             dependencies=dependencies,
             owner_class=analysis.owner_class,
+            closure_vars=closure_vars,
+            closure_func_refs=closure_func_refs,
             module_vars=analysis.module_vars,
         )
         self._nodes[qualified_name] = node
         self._funcs[qualified_name] = func
         logger.info("Auto-registered untraced dependency %s", qualified_name)
 
+        # Auto-register closure function deps (after node is in self._nodes to prevent re-entry)
+        for ref_qname, func_obj in closure_func_objects.items():
+            if ref_qname not in self._nodes:
+                self._auto_register(func_obj)
+
+        for cls_obj in closure_classes.values():
+            self._auto_register_class(cls_obj)
+
         self._discover_untraced_deps(func.__module__, node)
         return True
 

pyfuse/graph/store.py

@@ -1,5 +1,6 @@
 """Content-addressable store for serializing and reconstructing functions."""
 
+import ast
 import json
 import logging
 from typing import Any, Self
@@ -72,6 +73,46 @@ def _apply_closure_transforms(
     return source
 
 
+class _SuperRewriter(ast.NodeTransformer):
+    """Replace ``super()`` with ``super(ClassName, self)`` or ``super(ClassName, cls)``."""
+
+    def __init__(self, class_name: str) -> None:
+        self._class_name = class_name
+        self._first_param: str | None = None
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
+        self._first_param = node.args.args[0].arg if node.args.args else "self"
+        self.generic_visit(node)
+        self._first_param = None
+        return node
+
+    visit_AsyncFunctionDef = visit_FunctionDef  # type: ignore[assignment]
+
+    def visit_Call(self, node: ast.Call) -> ast.Call:
+        self.generic_visit(node)
+        if (
+            isinstance(node.func, ast.Name)
+            and node.func.id == "super"
+            and not node.args
+            and not node.keywords
+        ):
+            node.args = [
+                ast.Name(id=self._class_name, ctx=ast.Load()),
+                ast.Name(id=self._first_param or "self", ctx=ast.Load()),
+            ]
+        return node
+
+
+def _rewrite_bare_super(source: str, class_name: str) -> str:
+    """Replace zero-arg ``super()`` with ``super(ClassName, self/cls)``."""
+    if "super()" not in source:
+        return source
+    tree = ast.parse(source)
+    tree = _SuperRewriter(class_name).visit(tree)
+    ast.fix_missing_locations(tree)
+    return ast.unparse(tree)
+
+
 def _indent_method(source: str) -> str:
     """Indent a method source for embedding inside a class block."""
     return "\n".join(
@@ -89,7 +130,9 @@ def _build_class_block(
     class_name = owner_class.rsplit(".", 1)[-1]
 
     method_sources = [
-        _indent_method(_apply_closure_transforms(nodes[qname], nodes))
+        _indent_method(_rewrite_bare_super(
+            _apply_closure_transforms(nodes[qname], nodes), class_name,
+        ))
         for qname in member_qnames
     ]
 
@@ -119,7 +162,8 @@ def _build_class_block(
 
     decorator_lines = "".join(f"@{d}\n" for d in class_decorators)
     attr_block = "".join(
-        _indent_method(attr) + "\n\n" for attr in class_attrs
+        _indent_method(_rewrite_bare_super(attr, class_name)) + "\n\n"
+        for attr in class_attrs
     )
 
     return decorator_lines + header + attr_block + "\n\n".join(method_sources)

pyfuse/typing.py

@@ -36,9 +36,9 @@ async def run_in(self, delay: timedelta | float, *args: P.args, **kwargs: P.kwar
     async def run_every(
         self,
         frequency: timedelta | float,
-        *args: P.args,
+        *args: Any,
         _start_at: datetime | None = ...,
-        **kwargs: P.kwargs,
+        **kwargs: Any,
     ) -> ScheduleHandle: ...