Dataframe compatibility for transformers

src/sknnr/_base.py

@@ -1,12 +1,11 @@
 import warnings
-from contextlib import contextmanager
-from typing import Literal
 
 import numpy as np
-from sklearn.base import BaseEstimator
 from sklearn.neighbors import KNeighborsRegressor
 from sklearn.utils.validation import _get_feature_names, check_is_fitted
 
+from .transformers._base import set_temp_output
+
 
 class IDNeighborsRegressor(KNeighborsRegressor):
     """
@@ -80,15 +79,3 @@ def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
         return super().kneighbors(
             X=X_transformed, n_neighbors=n_neighbors, return_distance=return_distance
         )
-
-
-@contextmanager
-def set_temp_output(estimator: BaseEstimator, temp_mode: Literal["default", "pandas"]):
-    """Temporarily set the output mode of an estimator."""
-    previous_config = getattr(estimator, "_sklearn_output_config", {}).copy()
-
-    estimator.set_output(transform=temp_mode)
-    try:
-        yield
-    finally:
-        estimator._sklearn_output_config = previous_config

src/sknnr/transformers/_base.py

@@ -1,4 +1,8 @@
+from contextlib import contextmanager
+from typing import Literal
+
 import numpy as np
+from sklearn.base import TransformerMixin
 from sklearn.preprocessing import StandardScaler
 
 
@@ -11,3 +15,17 @@ def fit(self, X, y=None, sample_weight=None):
         scaler = super().fit(X, y, sample_weight)
         scaler.scale_ = np.std(np.asarray(X), axis=0, ddof=self.ddof)
         return scaler
+
+
+@contextmanager
+def set_temp_output(
+    transformer: TransformerMixin, temp_mode: Literal["default", "pandas"]
+):
+    """Temporarily set the output mode of an transformer."""
+    previous_config = getattr(transformer, "_sklearn_output_config", {}).copy()
+
+    transformer.set_output(transform=temp_mode)
+    try:
+        yield
+    finally:
+        transformer._sklearn_output_config = previous_config

src/sknnr/transformers/_ccora_transformer.py

@@ -1,5 +1,6 @@
 import numpy as np
 from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.utils.validation import check_is_fitted
 
 from . import StandardScalerWithDOF
 from ._ccora import CCorA
@@ -24,6 +25,7 @@ def fit(self, X, y):
         return self
 
     def transform(self, X, y=None):
+        check_is_fitted(self)
         return self.scaler_.transform(X) @ self.ccora_.projector
 
     def fit_transform(self, X, y):

src/sknnr/transformers/_mahalanobis_transformer.py

@@ -1,5 +1,6 @@
 import numpy as np
 from sklearn.base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
+from sklearn.utils.validation import check_is_fitted
 
 from . import StandardScalerWithDOF
 
@@ -14,6 +15,7 @@ def fit(self, X, y=None):
         return self
 
     def transform(self, X, y=None):
+        check_is_fitted(self)
         return self.scaler_.transform(X) @ self.transform_
 
     def fit_transform(self, X, y=None):

tests/test_transformers.py

@@ -111,7 +111,7 @@ def test_set_temp_output(moscow_euclidean, config_type):
         transformer.set_output(transform="pandas")
 
     # Temp output mode should override previously set config
-    with set_temp_output(estimator=transformer, temp_mode="default"):
+    with set_temp_output(transformer=transformer, temp_mode="default"):
         assert isinstance(transformer.transform(moscow_euclidean.X), np.ndarray)
 
     # Previous config should be restored