fix: scitype-correct export_code examples, is_pipeline, dataset, loaded models (#534)

Fixes #534.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Shashankss1205

src/sktime_mcp/tools/codegen.py

@@ -44,6 +44,86 @@ def _is_valid_var_name(var_name: str) -> bool:
     return isinstance(var_name, str) and var_name.isidentifier() and not keyword.iskeyword(var_name)
 
 
+def _loader_for(dataset: str, demo_datasets: dict) -> tuple[str, str]:
+    """Return (module, func) for a demo dataset name, defaulting to load_airline."""
+    if dataset in demo_datasets:
+        module_path = demo_datasets[dataset]
+        module, func = module_path.rsplit(".", 1)
+        return module, func
+    return "sktime.datasets", "load_airline"
+
+
+def _fit_example(
+    var_name: str,
+    obj_type: str,
+    dataset: str | None,
+    handle_info: Any,
+    demo_datasets: dict,
+) -> str:
+    """Build a runnable fit/predict example matching the estimator's scitype.
+
+    A forecaster-shaped example (`fit(y)` / `predict(fh)`) is wrong for
+    transformers, splitters, and classifiers, which raise AttributeError when
+    the generated code runs (BUG-03).
+    """
+    if obj_type in ("classifier", "regressor"):
+        # Panel X + label/target y — use a classification demo dataset.
+        ds = dataset or "arrow_head"
+        module, func = _loader_for(ds, demo_datasets)
+        verb = "class" if obj_type == "classifier" else "value"
+        return f"""
+
+# Example usage ({obj_type}):
+from {module} import {func}
+X, y = {func}(return_X_y=True)
+
+{var_name}.fit(X, y)
+predictions = {var_name}.predict(X)  # predicted {verb} per instance
+print(predictions)
+"""
+
+    if obj_type == "transformer":
+        ds = dataset or handle_info.metadata.get("training_dataset") or "airline"
+        module, func = _loader_for(ds, demo_datasets)
+        return f"""
+
+# Example usage (transformer):
+from {module} import {func}
+y = {func}()
+
+y_transformed = {var_name}.fit_transform(y)
+print(y_transformed)
+"""
+
+    if obj_type == "splitter":
+        ds = dataset or "airline"
+        module, func = _loader_for(ds, demo_datasets)
+        return f"""
+
+# Example usage (splitter):
+from {module} import {func}
+y = {func}()
+
+for train_idx, test_idx in {var_name}.split(y):
+    print("train:", train_idx, "test:", test_idx)
+"""
+
+    # Default: forecaster.
+    ds = dataset or handle_info.metadata.get("training_dataset") or "airline"
+    module, func = _loader_for(ds, demo_datasets)
+    return f"""
+
+# Example usage (forecaster):
+from {module} import {func}
+y = {func}()
+
+{var_name}.fit(y)
+fh = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]  # 12-step ahead forecast
+predictions = {var_name}.predict(fh=fh)
+print(predictions)
+"""
+
+
 def export_code_tool(
     handle: str,
     var_name: str = "model",
@@ -96,45 +176,42 @@ def export_code_tool(
 
     estimator_name = handle_info.estimator_name
     params = handle_info.params
-
     spec = params.get("spec")
-    if not spec:
-        return {"success": False, "error": "No craft spec found in handle parameters."}
 
-    is_pipeline = "*" in spec or "Pipeline" in spec or "[" in spec
-    code = f"from sktime.registry import craft\n\n{var_name} = craft({_format_value(spec)})"
+    instance = handle_manager.get_instance(handle)
+    get_tag = getattr(instance, "get_class_tag", None)
+    obj_type = get_tag("object_type", "") if callable(get_tag) else ""
+
+    # is_pipeline from the instance, not a spec substring — "[" in spec
+    # false-positived on any list argument (BUG-04).
+    is_pipeline = bool(spec and "*" in spec) or hasattr(instance, "steps")
+
+    if spec:
+        code = f"from sktime.registry import craft\n\n{var_name} = craft({_format_value(spec)})"
+    elif handle_info.metadata.get("source") == "loaded" and handle_info.metadata.get("path"):
+        # Loaded models carry no craft spec; emit a load_model snippet instead of
+        # failing with "No craft spec found" (NB-17).
+        model_path = handle_info.metadata["path"]
+        code = (
+            "from sktime.utils.mlflow_sktime import load_model\n\n"
+            f"{var_name} = load_model({_format_value(model_path)})"
+        )
+    else:
+        return {"success": False, "error": "No craft spec found in handle parameters."}
 
-    # Optionally add fit/predict example
+    # Optionally add a scitype-appropriate fit example (BUG-03).
     if include_fit_example:
-        # Priority: explicit argument > dataset used during fit > "airline" fallback
-        effective_dataset = dataset or handle_info.metadata.get("training_dataset") or "airline"
-        # Resolve the dataset loader from discovered demo datasets
         demo_datasets = _get_demo_datasets()
-        if effective_dataset in demo_datasets:
-            module_path = demo_datasets[effective_dataset]
-            module_parts = module_path.rsplit(".", 1)
-            loader_module = module_parts[0]
-            loader_func = module_parts[1]
-        else:
-            loader_module = "sktime.datasets"
-            loader_func = "load_airline"
-
-        example_code = f"""
-
-# Example usage:
-# Load data
-from {loader_module} import {loader_func}
-y = {loader_func}()
-
-# Fit the model
-{var_name}.fit(y)
-
-# Make predictions
-fh = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]  # 12-step ahead forecast
-predictions = {var_name}.predict(fh=fh)
-print(predictions)
-"""
-        code += example_code
+        if dataset is not None and dataset not in demo_datasets:
+            return {
+                "success": False,
+                "error": (
+                    f"Unknown dataset '{dataset}' for the fit example. Use a demo dataset "
+                    "name (see list_available_data) or omit dataset to use a default."
+                ),
+            }
+        example = _fit_example(var_name, obj_type, dataset, handle_info, demo_datasets)
+        code += example
 
     return {
         "success": True,

tests/test_export_code_scitype.py

@@ -0,0 +1,115 @@
+"""export_code must generate correct, runnable code per scitype (#534).
+
+- BUG-03: fit example was forecaster-shaped for every scitype (AttributeError
+  when run for transformers/splitters/classifiers).
+- BUG-04: is_pipeline false-positived on any "[" in the spec.
+- NB-15: the dataset argument was ignored; unknown datasets silently fell back
+  to airline.
+- NB-17: loaded models (empty params) failed with "No craft spec found".
+"""
+
+import contextlib
+
+import pytest
+
+from sktime_mcp.runtime.handles import get_handle_manager
+from sktime_mcp.tools.codegen import export_code_tool
+from sktime_mcp.tools.instantiate import instantiate_tool
+
+
+def _handle(spec):
+    res = instantiate_tool(spec=spec)
+    assert res["success"], res
+    return res["handle"]
+
+
+def _release(handle):
+    with contextlib.suppress(KeyError):
+        get_handle_manager().release_handle(handle)
+
+
+class TestIsPipeline:
+    def test_list_arg_is_not_a_pipeline(self):
+        h = _handle("SlidingWindowSplitter(window_length=24, fh=[1, 2, 3], step_length=12)")
+        try:
+            res = export_code_tool(h)
+            assert res["success"]
+            assert res["is_pipeline"] is False
+        finally:
+            _release(h)
+
+    def test_star_spec_is_a_pipeline(self):
+        h = _handle("Deseasonalizer() * NaiveForecaster()")
+        try:
+            res = export_code_tool(h)
+            assert res["success"]
+            assert res["is_pipeline"] is True
+        finally:
+            _release(h)
+
+
+class TestScitypeExampleRuns:
+    def _export_and_exec(self, spec):
+        h = _handle(spec)
+        try:
+            res = export_code_tool(h, include_fit_example=True)
+            assert res["success"], res
+            compile(res["code"], "<export>", "exec")
+            exec(res["code"], {})  # must run without AttributeError
+        finally:
+            _release(h)
+
+    def test_forecaster_example_runs(self):
+        self._export_and_exec("NaiveForecaster(sp=12)")
+
+    def test_transformer_example_runs(self):
+        self._export_and_exec("Deseasonalizer()")
+
+    def test_splitter_example_runs(self):
+        self._export_and_exec("SlidingWindowSplitter(window_length=24, step_length=12)")
+
+    def test_classifier_example_runs(self):
+        self._export_and_exec("KNeighborsTimeSeriesClassifier()")
+
+
+class TestLoadedModelExport:
+    def test_loaded_model_emits_load_model_snippet(self):
+        from sktime.forecasting.naive import NaiveForecaster
+
+        hm = get_handle_manager()
+        # Simulate a load_model handle: no craft spec, metadata carries the path.
+        handle = hm.create_handle(
+            estimator_name="NaiveForecaster",
+            instance=NaiveForecaster(),
+            params={},
+            metadata={"source": "loaded", "path": "/tmp/some_model_dir"},
+        )
+        try:
+            res = export_code_tool(handle)
+            assert res["success"], res
+            assert "load_model" in res["code"]
+            assert "/tmp/some_model_dir" in res["code"]
+            compile(res["code"], "<export>", "exec")
+        finally:
+            _release(handle)
+
+
+class TestDatasetValidation:
+    def test_unknown_dataset_rejected(self):
+        h = _handle("NaiveForecaster()")
+        try:
+            res = export_code_tool(h, include_fit_example=True, dataset="not_a_dataset_zzz")
+            assert not res["success"]
+            assert "not_a_dataset_zzz" in res["error"]
+        finally:
+            _release(h)
+
+    def test_known_dataset_used(self):
+        h = _handle("NaiveForecaster()")
+        try:
+            res = export_code_tool(h, include_fit_example=True, dataset="lynx")
+            assert res["success"]
+            assert "load_lynx" in res["code"]
+            assert "load_airline" not in res["code"]
+        finally:
+            _release(h)