chore: import upstream snapshot with attribution

tests/pyfunc/test_inferred_code_path.py

@@ -0,0 +1,149 @@
+import os
+import pickle
+
+import numpy as np
+import pytest
+import sklearn.datasets
+import sklearn.neighbors
+
+import mlflow
+from mlflow.models import Model
+
+
+@pytest.fixture
+def model_path(tmp_path):
+    return tmp_path / "model"
+
+
+@pytest.fixture(scope="module")
+def iris_data():
+    iris = sklearn.datasets.load_iris()
+    x = iris.data[:, :2]
+    y = iris.target
+    return x, y
+
+
+@pytest.fixture(scope="module")
+def sklearn_knn_model(iris_data):
+    x, y = iris_data
+    knn_model = sklearn.neighbors.KNeighborsClassifier()
+    knn_model.fit(x, y)
+    return knn_model
+
+
+def _walk_dir(path):
+    return {
+        str(p.relative_to(path))
+        for p in path.rglob("*")
+        if p.is_file() and p.parent.name != "__pycache__"
+    }
+
+
+def test_loader_module_model_save_load(
+    sklearn_knn_model, iris_data, tmp_path, model_path, monkeypatch
+):
+    monkeypatch.chdir(os.path.dirname(__file__))
+    monkeypatch.syspath_prepend(".")
+    sk_model_path = tmp_path / "knn.pkl"
+    with open(sk_model_path, "wb") as f:
+        pickle.dump(sklearn_knn_model, f)
+
+    model_config = Model(run_id="test", artifact_path="testtest")
+    mlflow.pyfunc.save_model(
+        path=model_path,
+        data_path=sk_model_path,
+        loader_module="custom_model.loader",
+        mlflow_model=model_config,
+        infer_code_paths=True,
+    )
+
+    reloaded_model_config = Model.load(model_path / "MLmodel")
+
+    assert _walk_dir(model_path / "code") == {
+        "custom_model/loader.py",
+        "custom_model/mod1/__init__.py",
+        "custom_model/mod1/mod2/__init__.py",
+        "custom_model/mod1/mod4.py",
+    }
+    assert model_config.__dict__ == reloaded_model_config.__dict__
+    assert mlflow.pyfunc.FLAVOR_NAME in reloaded_model_config.flavors
+    assert mlflow.pyfunc.PY_VERSION in reloaded_model_config.flavors[mlflow.pyfunc.FLAVOR_NAME]
+    reloaded_model = mlflow.pyfunc.load_model(model_path)
+    np.testing.assert_array_equal(
+        sklearn_knn_model.predict(iris_data[0]), reloaded_model.predict(iris_data[0])
+    )
+
+
+def get_model_class():
+    """
+    Defines a custom Python model class that wraps a scikit-learn estimator.
+    This can be invoked within a pytest fixture to define the class in the ``__main__`` scope.
+    Alternatively, it can be invoked within a module to define the class in the module's scope.
+    """
+    from custom_model.mod1 import mod2
+
+    class CustomSklearnModel(mlflow.pyfunc.PythonModel):
+        def __init__(self):
+            self.mod2 = mod2
+
+        def predict(self, context, model_input, params=None):
+            return [x + 10 for x in model_input]
+
+    return CustomSklearnModel
+
+
+def test_python_model_save_load(tmp_path, monkeypatch):
+    monkeypatch.chdir(os.path.dirname(__file__))
+    monkeypatch.syspath_prepend(".")
+
+    model_class = get_model_class()
+
+    pyfunc_model_path = tmp_path / "pyfunc_model"
+
+    mlflow.pyfunc.save_model(
+        path=pyfunc_model_path,
+        python_model=model_class(),
+        infer_code_paths=True,
+    )
+
+    assert _walk_dir(pyfunc_model_path / "code") == {
+        "custom_model/mod1/__init__.py",
+        "custom_model/mod1/mod2/__init__.py",
+        "custom_model/mod1/mod4.py",
+    }
+    loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri=pyfunc_model_path)
+    np.testing.assert_array_equal(
+        loaded_pyfunc_model.predict([1, 2, 3]),
+        [11, 12, 13],
+    )
+
+
+def test_transitive_import_capture(tmp_path, monkeypatch):
+    monkeypatch.chdir(os.path.dirname(__file__))
+    monkeypatch.syspath_prepend(".")
+
+    from custom_model.transitive_test.model_with_transitive import (
+        ModelWithTransitiveDependency,
+    )
+
+    pyfunc_model_path = tmp_path / "pyfunc_model"
+
+    mlflow.pyfunc.save_model(
+        path=pyfunc_model_path,
+        python_model=ModelWithTransitiveDependency(),
+        infer_code_paths=True,
+    )
+
+    # Verify that transitive_dependency.py is captured correctly
+    # This file is imported as "from ... import some_function" (importing a function)
+    # The fix ensures that we record the parent module when the imported item is not a module
+    assert _walk_dir(pyfunc_model_path / "code") == {
+        "custom_model/transitive_test/__init__.py",
+        "custom_model/transitive_test/model_with_transitive.py",
+        "custom_model/transitive_test/transitive_dependency.py",
+    }
+
+    # Verify the model works after loading
+    loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri=pyfunc_model_path)
+    result = loaded_pyfunc_model.predict([1, 2, 3])
+    assert result == ["test", "test", "test"]