diff --git a/scripts/detect_affected_models.py b/scripts/detect_affected_models.py index 04706162..4fb454f4 100644 --- a/scripts/detect_affected_models.py +++ b/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 diff --git a/scripts/detect_affected_models_test.py b/scripts/detect_affected_models_test.py index f9e07237..955df24e 100644 --- a/scripts/detect_affected_models_test.py +++ b/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,160 @@ 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. + monkeypatch.setattr("detect_affected_models._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(