CI: Narrow affected-models detection by resolving __init__.py re-exports

scripts/detect_affected_models.py

@@ -48,15 +48,15 @@
     "src/mobius/models/__init__.py",
 )
 
-# Task files are resolved by string-based lookup at runtime (not Python
-# imports), so the import graph cannot trace task → model dependencies.
-# Keep tasks/ as shared_infra until a task→model_type mapping exists.
-_SHARED_INFRA_PREFIXES = ("src/mobius/tasks/",)
+_SHARED_INFRA_PREFIXES: tuple[str, ...] = ()
 
 # Traceable infrastructure: component files that are analyzed via the
 # import graph to find which models they actually affect, rather than
 # triggering run_all unconditionally.
-_TRACEABLE_PREFIXES = ("src/mobius/components/",)
+_TRACEABLE_PREFIXES = (
+    "src/mobius/components/",
+    "src/mobius/tasks/",
+)
 
 
 def classify_file(path: str) -> str:
@@ -108,12 +108,24 @@ def classify_file(path: str) -> str:
 # ----------------------------------------------------------------
 
 
-def _parse_imports(filepath: Path) -> set[str]:
+def _parse_imports(
+    filepath: Path,
+    reexport_map: dict[tuple[str, str], str] | None = None,
+) -> set[str]:
     """Extract imported module names from a Python file using AST.
 
     Returns a set of dotted module names that appear in import
     statements. Only collects imports from within the
     mobius package.
+
+    When ``reexport_map`` is provided, ``from pkg import sym`` statements
+    are resolved through the re-export map to the underlying source
+    module that defines ``sym``. This avoids spurious dependencies on
+    re-export hubs like ``mobius.components/__init__.py``: a model that
+    imports ``Attention`` from ``mobius.components`` is recorded as
+    depending on ``mobius.components._attention`` (the actual source),
+    not on the package itself. Symbols not found in the re-export map
+    fall back to recording the package name.
     """
     try:
         source = filepath.read_text(encoding="utf-8")
@@ -129,10 +141,77 @@ def _parse_imports(filepath: Path) -> set[str]:
                     imports.add(alias.name)
         elif isinstance(node, ast.ImportFrom):
             if node.module and node.module.startswith("mobius"):
-                imports.add(node.module)
+                unresolved = False
+                for alias in node.names:
+                    if alias.name == "*":
+                        # Wildcard imports can't be resolved — fall back
+                        # to depending on the package itself.
+                        unresolved = True
+                        continue
+                    source_mod = (
+                        reexport_map.get((node.module, alias.name))
+                        if reexport_map is not None
+                        else None
+                    )
+                    if source_mod:
+                        imports.add(source_mod)
+                    else:
+                        unresolved = True
+                # Only record the package itself when at least one
+                # imported symbol could not be resolved through the
+                # re-export map. This avoids spurious dependencies on
+                # re-export hubs like ``mobius.components/__init__.py``.
+                if unresolved:
+                    imports.add(node.module)
     return imports
 
 
+def _build_reexport_map(search_dir: Path) -> dict[tuple[str, str], str]:
+    """Build a (package, symbol) → source_module map from ``__init__.py`` files.
+
+    Parses each ``__init__.py`` in the source tree for ``from .submodule
+    import Symbol`` and ``from mobius.pkg.submodule import Symbol``
+    statements. The resulting map lets us resolve re-exported symbols
+    back to their defining module so changes to a re-export hub don't
+    spuriously invalidate every importer.
+    """
+    reexport: dict[tuple[str, str], str] = {}
+    for init_file in search_dir.rglob("__init__.py"):
+        package = _module_name_from_path(init_file)
+        if not package:
+            continue
+        try:
+            source = init_file.read_text(encoding="utf-8")
+            tree = ast.parse(source, filename=str(init_file))
+        except (SyntaxError, UnicodeDecodeError):
+            continue
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.ImportFrom):
+                continue
+            # Resolve relative imports like ``from . import x`` or
+            # ``from .sub import X`` against the current package.
+            if node.level:
+                base_parts = package.split(".") if package else []
+                # ``from .`` keeps us at the same package; ``from ..`` goes up.
+                if node.level - 1 > len(base_parts):
+                    continue
+                base = ".".join(base_parts[: len(base_parts) - (node.level - 1)])
+                if node.module:
+                    src_module = f"{base}.{node.module}" if base else node.module
+                else:
+                    src_module = base
+            else:
+                src_module = node.module or ""
+            if not src_module.startswith("mobius"):
+                continue
+            for alias in node.names:
+                if alias.name == "*":
+                    continue
+                exported_name = alias.asname or alias.name
+                reexport[(package, exported_name)] = src_module
+    return reexport
+
+
 def _module_name_from_path(filepath: Path) -> str | None:
     """Convert a file path to a dotted module name.
 
@@ -159,12 +238,13 @@ def _build_import_graph(
     modules it directly imports.
     """
     graph: dict[str, set[str]] = {}
+    reexport_map = _build_reexport_map(search_dir)
     for pyfile in search_dir.rglob("*.py"):
         if pyfile.name.endswith("_test.py"):
             continue
         mod_name = _module_name_from_path(pyfile)
         if mod_name:
-            graph[mod_name] = _parse_imports(pyfile)
+            graph[mod_name] = _parse_imports(pyfile, reexport_map)
     return graph
 
 

scripts/detect_affected_models_test.py

@@ -22,9 +22,11 @@
     _SRC_ROOT,
     _build_class_to_source_module,
     _build_import_graph,
+    _build_reexport_map,
     _build_registry_class_to_types,
     _build_source_module_to_types,
     _find_reverse_dependents,
+    _parse_imports,
     classify_file,
     detect_affected_models,
 )
@@ -46,7 +48,7 @@ def test_component_file(self):
         assert classify_file("src/mobius/components/_attention.py") == "traceable"
 
     def test_task_file(self):
-        assert classify_file("src/mobius/tasks/_causal_lm.py") == "shared_infra"
+        assert classify_file("src/mobius/tasks/_causal_lm.py") == "traceable"
 
     def test_configs_file(self):
         assert classify_file("src/mobius/_configs.py") == "shared_infra"
@@ -181,10 +183,16 @@ def test_component_change_traces_affected_models(self):
         # _attention.py is imported by many models — should find affected types
         assert len(result["affected"]) > 0
 
-    def test_task_change_triggers_run_all(self):
-        """Task files use string-based lookup, not imports — must trigger run_all."""
+    def test_task_change_does_not_trigger_run_all(self):
+        """Task files are traceable but produce an empty affected set.
+
+        No model imports ``mobius.tasks`` directly (tasks are looked up at
+        runtime by string keys), so tracing through the import graph finds
+        no dependents. Documented limitation — see PR description.
+        """
         result = detect_affected_models(["src/mobius/tasks/_causal_lm.py"])
-        assert result["run_all"] is True
+        assert result["run_all"] is False
+        assert result["affected"] == []
 
     def test_configs_change_triggers_run_all(self):
         result = detect_affected_models(["src/mobius/_configs.py"])
@@ -354,10 +362,162 @@ def test_traceable_result_is_subset_of_all_models(self):
 
 
 # ----------------------------------------------------------------
-# CLI tests
+# Re-export resolution tests
+#
+# These tests use synthetic source trees in a temp directory so they
+# are isolated from the real mobius package layout.
 # ----------------------------------------------------------------
 
 
+class TestReexportResolution:
+    """Tests for _parse_imports + _build_reexport_map.
+
+    The resolver must record dependencies on the *source* module that
+    actually defines a symbol, not on re-export hubs like
+    ``components/__init__.py``. Wildcard and unknown symbols fall back
+    to depending on the hub package.
+    """
+
+    @staticmethod
+    def _write_pkg(tmp_path: Path, monkeypatch) -> Path:
+        """Create a synthetic ``src/mobius`` tree.
+
+        Layout::
+
+            src/mobius/__init__.py
+            src/mobius/components/__init__.py  # re-exports Attention, MLP
+            src/mobius/components/_attention.py  # defines Attention
+            src/mobius/components/_mlp.py        # defines MLP
+        """
+        src = tmp_path / "src" / "mobius"
+        (src / "components").mkdir(parents=True)
+        (src / "__init__.py").write_text("")
+        (src / "components" / "__init__.py").write_text(
+            "from ._attention import Attention\nfrom ._mlp import MLP\n"
+        )
+        (src / "components" / "_attention.py").write_text("class Attention: ...\n")
+        (src / "components" / "_mlp.py").write_text("class MLP: ...\n")
+        # _module_name_from_path uses _PROJECT_ROOT to resolve dotted names;
+        # point it at our synthetic tree for the duration of the test.
+        import detect_affected_models as _dam
+
+        monkeypatch.setattr(_dam, "_PROJECT_ROOT", tmp_path)
+        return src
+
+    def test_resolves_symbol_to_source_module(self, tmp_path: Path, monkeypatch) -> None:
+        """``from mobius.components import Attention`` → depends on _attention."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import Attention\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert "mobius.components._attention" in imports
+        # The hub package itself is NOT recorded when every symbol resolves.
+        assert "mobius.components" not in imports
+
+    def test_resolves_multiple_symbols_from_same_hub(
+        self, tmp_path: Path, monkeypatch
+    ) -> None:
+        """Each symbol in a multi-import resolves to its own source module."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import Attention, MLP\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert "mobius.components._attention" in imports
+        assert "mobius.components._mlp" in imports
+        assert "mobius.components" not in imports
+
+    def test_unknown_symbol_falls_back_to_package(self, tmp_path: Path, monkeypatch) -> None:
+        """Symbols not in the re-export map fall back to the hub package."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import NotExported\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        # Unknown symbol → conservative fallback on the package itself
+        assert "mobius.components" in imports
+        # And no spurious source-module resolution
+        assert "mobius.components._attention" not in imports
+
+    def test_wildcard_import_falls_back_to_package(self, tmp_path: Path, monkeypatch) -> None:
+        """``from mobius.components import *`` cannot be resolved — fall back."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import *\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert "mobius.components" in imports
+
+    def test_mixed_resolved_and_unresolved_records_both(
+        self, tmp_path: Path, monkeypatch
+    ) -> None:
+        """Mix of known + unknown symbols records resolved sources AND the hub."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import Attention, NotExported\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert "mobius.components._attention" in imports
+        # Unresolved symbol keeps the hub as a conservative dependency
+        assert "mobius.components" in imports
+
+    def test_plain_import_statement_records_module(self, tmp_path: Path, monkeypatch) -> None:
+        """``import mobius.components`` (no ``from``) records the module directly."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("import mobius.components\n")
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert "mobius.components" in imports
+
+    def test_no_reexport_map_records_module(self, tmp_path: Path, monkeypatch) -> None:
+        """When no map is passed, the hub is always recorded (legacy behavior)."""
+        self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text("from mobius.components import Attention\n")
+
+        imports = _parse_imports(importer, reexport_map=None)
+
+        assert "mobius.components" in imports
+        assert "mobius.components._attention" not in imports
+
+    def test_non_mobius_imports_ignored(self, tmp_path: Path, monkeypatch) -> None:
+        """Imports outside the mobius package are not recorded."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        importer = tmp_path / "importer.py"
+        importer.write_text(
+            "import os\nfrom typing import Any\nfrom mobius.components import Attention\n"
+        )
+
+        reexport = _build_reexport_map(src)
+        imports = _parse_imports(importer, reexport)
+
+        assert imports == {"mobius.components._attention"}
+
+    def test_reexport_map_built_from_relative_import(
+        self, tmp_path: Path, monkeypatch
+    ) -> None:
+        """``from .sub import X`` in __init__.py produces correct (pkg, X) entry."""
+        src = self._write_pkg(tmp_path, monkeypatch)
+        reexport = _build_reexport_map(src)
+
+        assert reexport[("mobius.components", "Attention")] == ("mobius.components._attention")
+        assert reexport[("mobius.components", "MLP")] == "mobius.components._mlp"
+
+
 class TestCLI:
     def test_json_output(self):
         result = subprocess.run(