Fix __init_subclass__

Co-authored-by: Copilot <copilot@github.com>

Author: codeSamuraii

docs/TECHNICAL_OVERVIEW.md

@@ -237,9 +237,11 @@ When a traced function calls an untraced user-defined function, pyfuse automatic
 
 Cross-module imports (e.g., `from utils import helper`) are converted from import statements to inline dependency edges, so the reconstructed code is self-contained.
 
-Class constructors (`MyClass()`) are auto-discovered: pyfuse registers all user-defined methods of the class. `@staticmethod` and `@classmethod` descriptors are unwrapped and registered correctly. When a method uses `super()`, base classes and their methods are discovered recursively, and `class Foo(Base):` headers are emitted in reconstructed source.
+Class constructors (`MyClass()`) are auto-discovered: pyfuse registers all user-defined methods of the class, even when the class relies on the implicit `object.__init__`. `@staticmethod` and `@classmethod` descriptors are unwrapped and registered correctly. Base classes and their methods are pulled in via `__mro__` walking (independent of `super()` usage), and `class Foo(Base):` headers are emitted in reconstructed source.
 
-Class-level attributes (assignments, annotated assignments, docstrings) are extracted from the class source AST and emitted in reconstructed class blocks. Class decorators (e.g., `@dataclass`) are captured and emitted above the class header. Metaclass keywords (e.g., `metaclass=ABCMeta`) and other class keywords are extracted from the class definition and included in the reconstructed header.
+Classes that hook `__init_subclass__` are treated as registries: every user-defined subclass is auto-registered so its definition fires the parent hook on the worker. The same path catches subclasses looked up indirectly (e.g., by name from a registry dict) without requiring an explicit reference in the traced body.
+
+Class-level attributes (assignments, annotated assignments, docstrings) are extracted from the class source AST and emitted in reconstructed class blocks. User classes referenced from class-body RHS expressions (e.g., descriptors installed via `field = Doubler()`) are auto-registered transitively. Class decorators (e.g., `@dataclass`) are captured and emitted above the class header. Metaclass keywords (e.g., `metaclass=ABCMeta`) and other class keywords are extracted from the class definition and included in the reconstructed header.
 
 Module-level constants and variables referenced by traced functions (e.g., `MAX_RETRIES = 5`) are captured and emitted in reconstructed source.
 
@@ -670,8 +672,7 @@ Environment variables `PYFUSE_SANDBOX_DOCKER_IMAGE` and `PYFUSE_SANDBOX_DOCKER_C
 - Aliased cross-module imports (`from utils import helper as h`) are skipped to avoid name mismatches.
 
 ### Classes
-- Metaclasses and `__init_subclass__` hooks are replayed when the parent class is in the dependency tree (i.e., referenced via `super()` or constructor call).
-- Class-level attributes defined via complex descriptors or external decorators (beyond simple assignments) may not be captured.
+- Class-level attributes are captured verbatim from the class source AST. User classes referenced from class-body expressions (e.g., descriptor instances assigned to class attributes) are auto-registered along with the owning class. External decorators applied to attributes via more dynamic patterns may still be missed.
 
 ## CLI reference
 

pyfuse/graph/analyzer.py

@@ -515,6 +515,16 @@ def _resolve_bare_calls(
             chosen = _prefer_same_module(init_matches, func_module)
             logger.debug("Constructor call %s() -> %s", called, chosen.qualified_name)
             deps.append(chosen.qualified_name)
+            continue
+        # Class with no user-defined ``__init__`` (inherits ``object.__init__``):
+        # link to every registered method so the whole class block is emitted.
+        method_matches = [
+            node for node in registry.values() if node.owner_class == called
+        ]
+        if method_matches:
+            for node in method_matches:
+                logger.debug("Bare class call %s() -> %s", called, node.qualified_name)
+                deps.append(node.qualified_name)
     return deps
 
 

pyfuse/graph/graph.py

@@ -291,6 +291,8 @@ def __init__(self) -> None:
         )
         self._runtime_deps: dict[str, set[str]] = {}
         self._lock: threading.Lock = threading.Lock()
+        self._classes_in_progress: set[str] = set()
+        self._inclusion_deps: dict[str, set[str]] = {}
 
     @classmethod
     def default(cls) -> "Graph":
@@ -499,23 +501,39 @@ def _auto_register_class(self, cls: type) -> None:
         """Auto-register all user-defined methods of a class into the graph."""
         class_name = cls.__name__
         module_name = cls.__module__
-
-        for attr_name, raw in cls.__dict__.items():
-            if isinstance(raw, (classmethod, staticmethod)):
-                func = raw.__func__
-            elif inspect.isfunction(raw):
-                func = raw
-            else:
-                continue
-            if not _is_user_function(func):
-                continue
-            qname = f"{module_name}.{class_name}.{attr_name}"
-            if qname in self._nodes:
-                continue
-            self._auto_register(func)
-
-        self._set_class_metadata(cls)
-        self._resolve_class_bases(cls)
+        cls_key = f"{module_name}.{cls.__qualname__}"
+        if cls_key in self._classes_in_progress:
+            return
+        self._classes_in_progress.add(cls_key)
+        try:
+            for attr_name, raw in cls.__dict__.items():
+                if isinstance(raw, (classmethod, staticmethod)):
+                    func = raw.__func__
+                elif inspect.isfunction(raw):
+                    func = raw
+                else:
+                    continue
+                if not _is_user_function(func):
+                    continue
+                qname = f"{module_name}.{class_name}.{attr_name}"
+                if qname in self._nodes:
+                    continue
+                self._auto_register(func)
+
+            self._set_class_metadata(cls)
+            self._resolve_class_bases(cls)
+
+            # Subclass registry pattern: classes that hook ``__init_subclass__``
+            # populate registries from subclass definitions. The traced source
+            # may look subclasses up indirectly (e.g. by name); to make that
+            # work on the worker, pull every user-defined subclass into the
+            # graph so its definition fires the parent hook on reconstruct.
+            if "__init_subclass__" in cls.__dict__:
+                for sub in cls.__subclasses__():
+                    if _is_user_class(sub):
+                        self._auto_register_class(sub)
+        finally:
+            self._classes_in_progress.discard(cls_key)
 
     def _set_class_metadata(self, cls: type) -> None:
         """Capture class-level attributes and decorators onto method nodes."""
@@ -532,10 +550,18 @@ def _set_class_metadata(self, cls: type) -> None:
         for deco_src in decorators:
             extra_names |= get_used_names(deco_src)
 
+        # User classes referenced from class-body RHS (e.g. descriptors like
+        # ``field = Doubler()``) are not visible to bare-call discovery on
+        # function bodies; register them here so they survive reconstruction.
+        ref_method_qnames = self._register_class_attr_refs(cls, extra_names)
+
         for node in self._nodes.values():
             if node.owner_class == class_name and node.module == module_name:
                 node.class_attrs = attrs
                 node.class_decorators = decorators
+                for ref_qname in ref_method_qnames:
+                    if ref_qname != node.qualified_name:
+                        self._inclusion_deps.setdefault(node.qualified_name, set()).add(ref_qname)
                 if extra_names:
                     existing_names = {imp.bound_name for imp in node.imports}
                     try:
@@ -551,6 +577,38 @@ def _set_class_metadata(self, cls: type) -> None:
                     except StopIteration:
                         pass
 
+    def _register_class_attr_refs(
+        self, cls: type, extra_names: set[str]
+    ) -> list[str]:
+        """Auto-register user classes referenced from the class body.
+
+        Returns the qualified names of one method per referenced class so
+        callers can wire dependency edges that keep them in the subgraph.
+        """
+        if not extra_names:
+            return []
+        module_obj = sys.modules.get(cls.__module__)
+        if module_obj is None:
+            return []
+
+        ref_method_qnames: list[str] = []
+        for name in extra_names:
+            if name in _BUILTIN_NAMES:
+                continue
+            obj = getattr(module_obj, name, None)
+            if obj is None or obj is cls:
+                continue
+            if not (inspect.isclass(obj) and _is_user_class(obj)):
+                continue
+            self._auto_register_class(obj)
+            for ref_node in self._nodes.values():
+                if (
+                    ref_node.owner_class == obj.__name__
+                    and ref_node.module == obj.__module__
+                ):
+                    ref_method_qnames.append(ref_node.qualified_name)
+        return ref_method_qnames
+
     def _resolve_class_bases(self, cls: type) -> None:
         """Detect class bases and store them on method nodes.
 
@@ -592,6 +650,36 @@ def _resolve_class_bases(self, cls: type) -> None:
             if _is_user_class(base_cls):
                 self._auto_register_class(base_cls)
 
+        # Add ordering edges from child methods to every parent method per
+        # direct user base, so the topological reconstruction emits parents
+        # first when the subclass is included via the registry pattern
+        # (without relying on ``super()`` being present in the subclass body).
+        for base_cls in cls.__bases__:
+            if base_cls is object or not _is_user_class(base_cls):
+                continue
+            parent_method_qnames = [
+                parent_node.qualified_name
+                for parent_node in self._nodes.values()
+                if (
+                    parent_node.owner_class == base_cls.__name__
+                    and parent_node.module == base_cls.__module__
+                )
+            ]
+            if not parent_method_qnames:
+                continue
+            for child_node in self._nodes.values():
+                if (
+                    child_node.owner_class != class_name
+                    or child_node.module != module_name
+                ):
+                    continue
+                bucket = self._inclusion_deps.setdefault(
+                    child_node.qualified_name, set()
+                )
+                for parent_qname in parent_method_qnames:
+                    if parent_qname != child_node.qualified_name:
+                        bucket.add(parent_qname)
+
     def _discover_untraced_deps(
         self, module_name: str, node: FunctionNode
     ) -> None:
@@ -609,6 +697,7 @@ def _discover_untraced_deps(
         self._discover_bare_call_deps(node, module_obj)
         if node.owner_class:
             self._discover_self_call_deps(node, module_obj, module_name)
+        self._discover_init_subclass_deps(node, module_obj)
 
     def _discover_bare_call_deps(
         self,
@@ -675,6 +764,51 @@ def _discover_self_call_deps(
             self._set_class_metadata(cls_obj)
             self._resolve_class_bases(cls_obj)
 
+    def _discover_init_subclass_deps(
+        self,
+        node: FunctionNode,
+        module_obj: object,
+    ) -> None:
+        """Pull subclasses of registry-style parents into the caller's deps.
+
+        When the traced source references a class with a user-defined
+        ``__init_subclass__``, its subclasses participate by being defined --
+        not by being named in the source.  Add inclusion edges from this
+        node to one method of each user subclass so the subgraph keeps them.
+        """
+        for name in find_bare_calls(node.source) | get_used_names(node.source):
+            if name in _BUILTIN_NAMES:
+                continue
+            obj = getattr(module_obj, name, None)
+            if obj is None or not inspect.isclass(obj):
+                continue
+            if "__init_subclass__" not in obj.__dict__:
+                continue
+            # Skip when this node is itself a method of obj or an ancestor;
+            # adding child-class edges from a parent method causes cycles
+            # via the parent ordering edges added in ``_resolve_class_bases``.
+            if node.owner_class is not None:
+                node_cls = getattr(module_obj, node.owner_class.rsplit(".", 1)[-1], None)
+                if (
+                    inspect.isclass(node_cls)
+                    and node_cls is not None
+                    and (node_cls is obj or issubclass(obj, node_cls))
+                ):
+                    continue
+            for sub in obj.__subclasses__():
+                if not _is_user_class(sub):
+                    continue
+                self._auto_register_class(sub)
+                for sub_node in self._nodes.values():
+                    if (
+                        sub_node.owner_class == sub.__name__
+                        and sub_node.module == sub.__module__
+                        and sub_node.qualified_name != node.qualified_name
+                    ):
+                        self._inclusion_deps.setdefault(
+                            node.qualified_name, set()
+                        ).add(sub_node.qualified_name)
+
     # -- Refresh & dependency merging ------------------------------------------
 
     def refresh(self) -> None:
@@ -693,6 +827,9 @@ def refresh(self) -> None:
             for ref_qname in node.closure_func_refs.values():
                 if ref_qname != qname and ref_qname not in deps:
                     deps.append(ref_qname)
+            for incl_qname in self._inclusion_deps.get(qname, set()):
+                if incl_qname in self._nodes and incl_qname not in deps:
+                    deps.append(incl_qname)
             node.dependencies = deps
 
     def _add_super_deps(self) -> None:

tests/test_auto_discovery.py

@@ -939,3 +939,113 @@ def test_class_docstring_captured(tmp_path: Path) -> None:
     )
     source = reconstruct(serialize(), "method")
     assert "A documented class" in source
+
+# ---------------------------------------------------------------------------
+# Feature 5: Bare class call without user-defined ``__init__``
+# ---------------------------------------------------------------------------
+
+
+def test_bare_class_no_init_chained(tmp_path: Path) -> None:
+    """``Tag().label()`` works when ``Tag`` has no user-defined ``__init__``."""
+    create_module(
+        tmp_path,
+        "adbarenoinit",
+        (
+            "from pyfuse import trace\n\n"
+            "class Tag:\n"
+            "    def label(self):\n"
+            "        return 'tag'\n\n"
+            "@trace\n"
+            "def construct_tag():\n"
+            "    return Tag().label()\n"
+        ),
+    )
+    source = reconstruct(serialize(), "construct_tag")
+    assert "class Tag" in source
+    assert "def label" in source
+
+    ns: dict[str, object] = {}
+    exec(source, ns)  # noqa: S102
+    assert ns["construct_tag"]() == "tag"  # type: ignore[operator]
+
+
+# ---------------------------------------------------------------------------
+# Feature 6: ``__init_subclass__`` registry pattern
+# ---------------------------------------------------------------------------
+
+
+def test_init_subclass_registry_pulls_subclasses(tmp_path: Path) -> None:
+    """Subclasses of registry-style parents survive serialization."""
+    create_module(
+        tmp_path,
+        "adsubregistry",
+        (
+            "from pyfuse import trace\n\n"
+            "class Plugin:\n"
+            "    _registry = {}\n"
+            "    def __init__(self):\n"
+            "        pass\n"
+            "    def __init_subclass__(cls, name=None, **kwargs):\n"
+            "        super().__init_subclass__(**kwargs)\n"
+            "        if name is not None:\n"
+            "            Plugin._registry[name] = cls\n"
+            "    def run(self, payload):\n"
+            "        raise NotImplementedError\n\n"
+            "class GreetPlugin(Plugin, name='greet'):\n"
+            "    def run(self, payload):\n"
+            "        return f'hello, {payload}!'\n\n"
+            "@trace\n"
+            "def dispatch(plugin_name, payload):\n"
+            "    Plugin()\n"
+            "    cls = Plugin._registry[plugin_name]\n"
+            "    return cls().run(payload)\n"
+        ),
+    )
+    source = reconstruct(serialize(), "dispatch")
+    assert "class Plugin" in source
+    assert "class GreetPlugin(Plugin, name='greet'):" in source
+
+    ns: dict[str, object] = {}
+    exec(source, ns)  # noqa: S102
+    assert ns["dispatch"]("greet", "world") == "hello, world!"  # type: ignore[operator]
+
+
+# ---------------------------------------------------------------------------
+# Feature 7: Class-level descriptor (user class referenced in class body)
+# ---------------------------------------------------------------------------
+
+
+def test_class_level_descriptor_registered(tmp_path: Path) -> None:
+    """A user class referenced from a class body (``field = Doubler()``) is captured."""
+    create_module(
+        tmp_path,
+        "addescriptor",
+        (
+            "from pyfuse import trace\n\n"
+            "class Doubler:\n"
+            "    def __set_name__(self, owner, name):\n"
+            "        self._attr = f'_{name}'\n"
+            "    def __get__(self, obj, objtype=None):\n"
+            "        if obj is None:\n"
+            "            return self\n"
+            "        return getattr(obj, self._attr) * 2\n"
+            "    def __set__(self, obj, value):\n"
+            "        setattr(obj, self._attr, value)\n\n"
+            "class Counter:\n"
+            "    value = Doubler()\n"
+            "    def __init__(self, start):\n"
+            "        self.value = start\n\n"
+            "@trace\n"
+            "def doubled(start):\n"
+            "    c = Counter(start)\n"
+            "    return c.value\n"
+        ),
+    )
+    source = reconstruct(serialize(), "doubled")
+    assert "class Doubler" in source
+    assert "class Counter" in source
+    assert "value = Doubler()" in source
+
+    ns: dict[str, object] = {}
+    exec(source, ns)  # noqa: S102
+    assert ns["doubled"](5) == 10  # type: ignore[operator]