diff --git a/skops/io/_trusted_types.py b/skops/io/_trusted_types.py index 1f337b95..a861145b 100644 --- a/skops/io/_trusted_types.py +++ b/skops/io/_trusted_types.py @@ -78,6 +78,14 @@ except ImportError: pass +try: + from sklearn.metrics._dist_metrics import EuclideanDistance64 + from sklearn.neighbors._kd_tree import KDTree + + _SKLEARN_INTERNAL_TYPES.extend([EuclideanDistance64, KDTree]) +except ImportError: + pass + SKLEARN_INTERNAL_TYPE_NAMES = [ get_type_name(t) for t in _SKLEARN_INTERNAL_TYPES diff --git a/skops/io/tests/test_persist.py b/skops/io/tests/test_persist.py index 656f9544..c89150f5 100644 --- a/skops/io/tests/test_persist.py +++ b/skops/io/tests/test_persist.py @@ -553,6 +553,29 @@ def test_gradient_boosting_estimators_have_no_untrusted_types(estimator, problem assert_method_outputs_equal(estimator, loaded, X) +def test_kneighbors_has_no_untrusted_types(): + """KNeighborsClassifier with kd_tree should not require manual trusted types.""" + from sklearn.neighbors import KNeighborsClassifier + + X, y = make_classification( + n_samples=N_SAMPLES, n_features=N_FEATURES, random_state=0 + ) + clf = KNeighborsClassifier(algorithm="kd_tree").fit(X, y) + + dumped = dumps(clf) + untrusted = get_untrusted_types(data=dumped) + knn_types = { + "sklearn.metrics._dist_metrics.EuclideanDistance64", + "sklearn.neighbors._kd_tree.KDTree", + } + assert not knn_types.intersection(untrusted), ( + f"KNN internal types still untrusted: {knn_types & set(untrusted)}" + ) + + loaded = loads(dumped) + assert_method_outputs_equal(clf, loaded, X) + + @pytest.mark.parametrize( "estimator", list(_unsupported_estimators()), ids=_get_check_estimator_ids )