Merge pull request #22 from yu9824/fix/support-sklearn-1_8

feat: support sklearn `ensure_all_finite` (Resolve #21)

Author: yu9824

src/kennard_stone/_core/_core.py

@@ -16,6 +16,8 @@
     from typing import Callable, Generator
 
 
+from inspect import signature
+
 import numpy as np
 from numpy.typing import ArrayLike
 from sklearn.feature_selection import VarianceThreshold
@@ -195,19 +197,6 @@ def _iter_indices(
             yield ind_train.tolist(), ind_test.tolist()
 
 
-@overload
-def train_test_split(
-    *arrays: T,
-    test_size: Optional[Union[float, int]] = None,
-    train_size: Optional[Union[float, int]] = None,
-    metric: Union[
-        Metrics, Callable[[ArrayLike, ArrayLike], np.ndarray]
-    ] = "euclidean",
-    n_jobs: Optional[int] = None,
-    device: Device = "cpu",
-) -> list[T]: ...
-
-
 def train_test_split(
     *arrays: T,
     test_size: Optional[Union[float, int]] = None,
@@ -426,14 +415,42 @@ def get_indexes(self, X: ArrayLike) -> list[array[int]]:
             The sorted indexes.
         """
         # check input array
-        X_checked: np.ndarray = check_array(
-            X,
+        # scikit-learn 1.6+ deprecates 'force_all_finite' and 1.8 renames to
+        # 'ensure_all_finite'. Check the signature dynamically.
+        check_array_sig = signature(check_array)
+        supports_ensure_all_finite = (
+            "ensure_all_finite" in check_array_sig.parameters
+        )
+        supports_force_all_finite = (
+            "force_all_finite" in check_array_sig.parameters
+        )
+
+        check_kwargs: dict[str, Any] = dict(
             ensure_2d=True,
             dtype="numeric",
-            force_all_finite="allow-nan"
-            if self.metric == "nan_euclidean"
-            else True,
         )
+        if supports_ensure_all_finite:
+            check_kwargs["ensure_all_finite"] = (
+                "allow-nan" if self.metric == "nan_euclidean" else True
+            )
+        elif supports_force_all_finite:
+            check_kwargs["force_all_finite"] = (
+                "allow-nan" if self.metric == "nan_euclidean" else True
+            )
+
+        try:
+            X_checked: np.ndarray = check_array(
+                X,
+                **check_kwargs,
+            )
+        except TypeError:
+            # Fallback when the argument is not accepted at runtime
+            check_kwargs.pop("ensure_all_finite", None)
+            check_kwargs.pop("force_all_finite", None)
+            X_checked = check_array(
+                X,
+                **check_kwargs,
+            )
         n_samples = X_checked.shape[0]
 
         # drop no variance

src/kennard_stone/utils/_pairwise.py

@@ -7,6 +7,8 @@
 else:
     from typing import Callable
 
+from inspect import signature
+
 import numpy as np
 import sklearn.metrics.pairwise
 from numpy.typing import ArrayLike
@@ -25,7 +27,8 @@ def pairwise_distances(
         Metrics, Callable[[ArrayLike, ArrayLike], np.ndarray]
     ] = "euclidean",
     n_jobs: Optional[int] = None,
-    force_all_finite=True,
+    ensure_all_finite: bool = True,
+    force_all_finite: Optional[bool] = None,
     device: Device = "cpu",
     verbose: int = 1,
     **kwargs,
@@ -60,9 +63,13 @@ def pairwise_distances(
         down the pairwise matrix into n_jobs even slices and computing them in
         parallel. (Note: 'n_jobs' is not supported by PyTorch.)
 
-    force_all_finite : bool, default=True
+    ensure_all_finite : bool, default=True
         Whether to raise an error on np.inf and np.nan in X.
 
+    force_all_finite : Optional[bool], default=None
+        Deprecated alias of 'ensure_all_finite'. If provided, a warning is
+        emitted and its value overrides 'ensure_all_finite'.
+
     device : Literal['cpu', 'cuda', 'mps'] or torch.device or str
     , default="cpu"
         Device to use for calculating pairwise distances.
@@ -80,6 +87,14 @@ def pairwise_distances(
     else:
         available_torch = False
 
+    # Handle deprecated alias
+    if force_all_finite is not None:
+        warnings.warn(
+            "'force_all_finite' is deprecated. Use 'ensure_all_finite' instead.",
+            DeprecationWarning,
+        )
+        ensure_all_finite = force_all_finite
+
     if available_torch:
         # Convert NumPy array to PyTorch tensor and move it to GPU
         X_tensor = torch.tensor(X, dtype=torch.float32, device=device)
@@ -109,14 +124,34 @@ def pairwise_distances(
             _logger.info(
                 "Calculating pairwise distances using scikit-learn.\n"
             )
-        return sklearn.metrics.pairwise.pairwise_distances(
-            X,
+        # scikit-learn 1.6+ deprecates 'force_all_finite' and 1.8 renames to
+        # 'ensure_all_finite'. Dynamically use whichever is available.
+        pd_sig = signature(sklearn.metrics.pairwise.pairwise_distances)
+        supports_ensure_all_finite = "ensure_all_finite" in pd_sig.parameters
+        supports_force_all_finite = "force_all_finite" in pd_sig.parameters
+
+        call_kwargs = dict(
             Y=Y,
             metric=metric,
             n_jobs=n_jobs,
-            force_all_finite=force_all_finite,
             **kwargs,
         )
+        if supports_ensure_all_finite:
+            call_kwargs["ensure_all_finite"] = ensure_all_finite
+        elif supports_force_all_finite:
+            call_kwargs["force_all_finite"] = ensure_all_finite
+
+        try:
+            return sklearn.metrics.pairwise.pairwise_distances(
+                X, **call_kwargs
+            )
+        except TypeError:
+            # Fallback for environments where the arg is rejected at runtime
+            call_kwargs.pop("ensure_all_finite", None)
+            call_kwargs.pop("force_all_finite", None)
+            return sklearn.metrics.pairwise.pairwise_distances(
+                X, **call_kwargs
+            )
     else:
         if verbose > 0:
             _logger.info(