get_feature_names_out for IdentityTransformer and test simplifica…

…tion.

sklego/preprocessing/identitytransformer.py

@@ -1,10 +1,10 @@
-from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.base import BaseEstimator, TransformerMixin, _OneToOneFeatureMixin
 
 from sklearn.utils import check_array
 from sklearn.utils.validation import check_is_fitted
 
 
-class IdentityTransformer(BaseEstimator, TransformerMixin):
+class IdentityTransformer(BaseEstimator, TransformerMixin, _OneToOneFeatureMixin):
     """
     The identity transformer returns what it is fed. Does not apply anything useful.
     The reason for having it is because you can build more expressive pipelines.
@@ -26,6 +26,7 @@ def fit(self, X, y=None):
         if self.check_X:
             X = check_array(X, copy=True, estimator=self)
         self.shape_ = X.shape
+        self.n_features_in_ = X.shape[1]
         return self
 
     def transform(self, X):

tests/test_preprocessing/test_columncapper.py

@@ -135,7 +135,7 @@ def test_get_feature_names_out(random_xy_dataset_clf):
     expected_feature_names = [f"feature_{i}" for i in n_col_range]
     np.testing.assert_array_equal(feature_names, expected_feature_names)
 
-    # Make sure get_feature_names_out cannot be called without given input_features if ColumnCapper is not fitted.
+    # get_feature_names_out should not work without given input_features if ColumnCapper is not fitted.
     with pytest.raises(ValueError):
         cc.get_feature_names_out(input_features=None)
 

tests/test_preprocessing/test_dictmapper.py

@@ -72,7 +72,7 @@ def test_get_feature_names_out(random_xy_dataset_clf):
     expected_feature_names = ['foobar_feature']
     np.testing.assert_array_equal(feature_names, expected_feature_names)
 
-    # Make sure get_feature_names_out cannot be called without given input_features if DictMapper is not fitted.
+    # get_feature_names_out should not work without given input_features if DictMapper is not fitted.
     with pytest.raises(ValueError):
         dm.get_feature_names_out(input_features=None)
 

tests/test_preprocessing/test_identitytransformer.py

@@ -36,3 +36,18 @@ def test_nan_inf(random_xy_dataset_regr):
     X[np.random.ranf(size=X.shape) > 0.9] = -np.inf
     X[np.random.ranf(size=X.shape) > 0.9] = np.inf
     X_new = IdentityTransformer(check_X=False).fit_transform(X)
+
+
+def test_get_feature_names_in(random_xy_dataset_regr):
+    X, y = random_xy_dataset_regr
+    it = IdentityTransformer()
+
+    # get_feature_names_out should not work without given input_features if IdentityTransformer is not fitted.
+    with pytest.raises(ValueError):
+        it.get_feature_names_out(input_features=None)
+
+    # Test with no input_features after being fitted
+    it.fit(X, y)
+    feature_names = it.get_feature_names_out()
+    expected_feature_names = [f"x{i}" for i in range(X.shape[1])]
+    np.testing.assert_array_equal(feature_names, expected_feature_names)