save_component carries the component's imports with it

src/ml4t_coursework/components.py

@@ -8,10 +8,14 @@
 
 from __future__ import annotations
 
+import ast
 import datetime as dt
+import importlib
 import inspect
 import json
+import sys
 import textwrap
+import types
 from collections.abc import Callable, Sequence
 from typing import Any
 
@@ -49,6 +53,102 @@ def _source_of(obj: Any, name: str) -> str:
         ) from exc
 
 
+# --- carrying the component's imports with it -------------------------------------------------
+#
+# A component is validated on its source re-executed in an empty namespace, so anything it reads
+# from another cell has to travel with it. `also` and `include` carry a function and a value. An
+# import is the third case and neither of those fits it: `also` would try to inline the library's
+# own source and `include` would write its repr into the file. So the import travels as what it
+# is, an import statement, worked out from the names the component actually mentions.
+
+
+def _mentioned(source: str) -> set[str]:
+    """Every bare name and attribute root the source refers to.
+
+    Deliberately not scope-aware. It over-collects - a local variable's name lands here too - and
+    the caller-namespace lookup below is what filters, because a local name does not resolve to an
+    importable object. Precise scoping would be more code for the same result.
+    """
+    try:
+        tree = ast.parse(textwrap.dedent(source))
+    except SyntaxError:
+        return set()
+    found: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load):
+            found.add(node.id)
+        elif isinstance(node, ast.Attribute):
+            root = node
+            while isinstance(root, ast.Attribute):
+                root = root.value
+            if isinstance(root, ast.Name):
+                found.add(root.id)
+    return found
+
+
+def _import_line(name: str, obj: Any) -> str | None:
+    """The import that brings `obj` back as `name` on a cold session, or None if there is none."""
+    if isinstance(obj, types.ModuleType):
+        module = obj.__name__
+        return f"import {module}" if module == name else f"import {module} as {name}"
+    module = getattr(obj, "__module__", None)
+    symbol = getattr(obj, "__qualname__", None) or getattr(obj, "__name__", None)
+    if not module or not symbol or "." in symbol or module in {"__main__", "builtins"}:
+        return None
+    module = _public_home(module, symbol, obj)
+    if module is None:
+        return None
+    return f"from {module} import {symbol}" if symbol == name else (
+        f"from {module} import {symbol} as {name}")
+
+
+def _public_home(module: str, symbol: str, obj: Any) -> str | None:
+    """The shortest importable path that re-exports `obj`, which is the one the student typed.
+
+    `Ridge.__module__` is `sklearn.linear_model._ridge`, and writing that into a student's file
+    records a private path that the library is free to rename. Walking up to the shallowest
+    ancestor that still exports the same object recovers `sklearn.linear_model`.
+    """
+    parts = module.split(".")
+    best = None
+    for depth in range(1, len(parts) + 1):
+        candidate = ".".join(parts[:depth])
+        try:
+            found = importlib.import_module(candidate)
+        except Exception:
+            continue
+        if getattr(found, symbol, None) is obj:
+            best = candidate
+            break
+    if best is None:
+        # Not reachable from its own package - a class defined in a notebook cell, most often.
+        return None
+    return best
+
+
+def _caller_namespace(depth: int) -> dict[str, Any]:
+    """The notebook cell's names, as seen from `depth` frames above this one."""
+    try:
+        frame = sys._getframe(depth)
+    except ValueError:  # pragma: no cover - only if the stack is shallower than the call
+        return {}
+    return {**frame.f_globals, **frame.f_locals}
+
+
+def _carried_imports(source: str, namespace: dict[str, Any]) -> list[str]:
+    """The import statements the saved file needs so it stands on its own."""
+    provided = {"np", "pd", "numpy", "pandas"}
+    lines = []
+    for name in sorted(_mentioned(source) - provided):
+        obj = namespace.get(name)
+        if obj is None:
+            continue
+        line = _import_line(name, obj)
+        if line:
+            lines.append(line)
+    return lines
+
+
 def _rebuild(source: str, symbol: str, name: str) -> Any:
     """Execute the saved source in a fresh namespace and return the object.
 
@@ -114,9 +214,11 @@ def save_component(
             raise ValueError(
                 f"{name}: pass the function or class itself, by name, not an instance or a lambda."
             )
-        parts = [f"{key} = {value!r}" for key, value in (include or {}).items()]
-        parts += [_source_of(helper, name) for helper in also]
-        parts.append(_source_of(obj, name))
+        written = [f"{key} = {value!r}" for key, value in (include or {}).items()]
+        written += [_source_of(helper, name) for helper in also]
+        written.append(_source_of(obj, name))
+        carried = _carried_imports("\n\n".join(written), _caller_namespace(2))
+        parts = (["\n".join(carried)] if carried else []) + written
         body = "\n\n".join(parts)
         checked = _rebuild(body, symbol, name)
         target = folder / f"{name}.py"

tests/test_helper.py

@@ -67,6 +67,75 @@ def fold_splitter(index, n_folds=4):
     assert "include=" in detail or "also=" in detail, "and say how to fix it"
 
 
+def test_an_imported_class_travels_with_the_component(capsys):
+    """The normal thing a student does: import in one cell, use it in the component.
+
+    Neither `also` nor `include` can carry an import - one inlines the library's source, the
+    other writes its repr - so the import has to be worked out and written as an import.
+    """
+    from sklearn.linear_model import Ridge
+
+    class LinearModel:
+        def __init__(self):
+            self.model = Ridge(alpha=1.0)
+            self.columns_ = None
+
+        def fit(self, X, y):
+            self.columns_ = list(X.columns)
+            self.model.fit(X, y)
+            return self
+
+        def predict(self, X):
+            if self.columns_ is None:
+                raise RuntimeError("fit the model before predicting")
+            return pd.Series(self.model.predict(X[self.columns_]), index=X.index,
+                             name="prediction")
+
+        @property
+        def coef_(self):
+            return self.model.coef_
+
+    result = save_component("model_linear", LinearModel, quiet=True)
+    assert result.conformant, [c.detail for c in result.checks if not c.passed]
+
+    saved = (project.components_dir() / "model_linear.py").read_text()
+    assert "from sklearn.linear_model import Ridge" in saved
+    assert "class Ridge" not in saved, "the import travels, not the library's source"
+
+    assert load_component("model_linear", quiet=True) is not None
+    assert source_of("model_linear") == "yours"
+
+
+def test_a_module_imported_under_an_alias_travels_too():
+    import numpy.linalg as la
+
+    def fold_splitter(index, n_folds=4):
+        index = pd.Index(index).sort_values()
+        assert la.norm([1.0]) == 1.0
+        block = len(index) // (n_folds + 1)
+        return [(index[: block * (k + 1) - 21], index[block * (k + 1): block * (k + 2)])
+                for k in range(n_folds)]
+
+    assert save_component("fold_splitter", fold_splitter, quiet=True).conformant
+    saved = (project.components_dir() / "fold_splitter.py").read_text()
+    assert "import numpy.linalg as la" in saved
+
+
+def test_a_value_from_another_cell_is_still_refused(capsys):
+    """Carrying imports must not start swallowing the genuine missing-symbol case."""
+    outside = 21
+
+    def fold_splitter(index, n_folds=4):
+        index = pd.Index(index).sort_values()
+        block = len(index) // (n_folds + 1)
+        return [(index[: block * (k + 1) - outside], index[block * (k + 1): block * (k + 2)])
+                for k in range(n_folds)]
+
+    result = save_component("fold_splitter", fold_splitter)
+    assert not result.conformant
+    assert "outside" in [c.detail for c in result.checks if not c.passed][0]
+
+
 def test_include_carries_a_value_the_component_reads():
     threshold = 21