feature: add get_feature_names() and tests to FunctionTransformer

nelson-liu · nelson-liu · commit a53f8e4801c3 · 2016-02-23T20:53:20.000-08:00
diff --git a/sklearn/preprocessing/_function_transformer.py b/sklearn/preprocessing/_function_transformer.py
@@ -64,11 +64,31 @@ def fit(self, X, y=None):
             check_array(X, self.accept_sparse)
         return self
 
+    def get_feature_names(self, input_features=None):
+        """
+        Return feature names for output features
+
+        Parameters
+        ----------
+        input_features : list of string, length len(input_features), optional
+            String names for input features if available. By default,
+            None is used.
+
+        Returns
+        -------
+        output_feature_names : list of string, length len(input_features)
+
+        """
+        if input_features is not None:
+            input_features = list(map(lambda input_feature: 'f(' +
+                                 input_feature + ')',
+                                 input_features))
+        return input_features
+
     def transform(self, X, y=None):
         if self.validate:
             X = check_array(X, self.accept_sparse)
         func = self.func if self.func is not None else _identity
 
-
         return func(X, *((y,) if self.pass_y else ()),
                     **(self.kw_args if self.kw_args else {}))
diff --git a/sklearn/preprocessing/tests/test_function_transformer.py b/sklearn/preprocessing/tests/test_function_transformer.py
@@ -74,6 +74,14 @@ def test_delegate_to_func():
     )
 
 
+def test_get_feature_names():
+    F = FunctionTransformer()
+    feature_names = F.get_feature_names(["a", "b", "c"])
+    testing.assert_array_equal(['f(a)', 'f(b)', 'f(c)'], feature_names)
+    feature_names_default = F.get_feature_names()
+    testing.assert_array_equal(None, feature_names_default)
+
+
 def test_np_log():
     X = np.arange(10).reshape((5, 2))
 
@@ -91,7 +99,7 @@ def test_kw_arg():
 
     # Test that rounding is correct
     testing.assert_array_equal(F.transform(X),
-                                  np.around(X, decimals=3))
+                               np.around(X, decimals=3))
 
 
 def test_kw_arg_update():
@@ -100,10 +108,9 @@ def test_kw_arg_update():
     F = FunctionTransformer(np.around, kw_args=dict(decimals=3))
 
     F.kw_args['decimals'] = 1
-
     # Test that rounding is correct
     testing.assert_array_equal(F.transform(X),
-                                  np.around(X, decimals=1))
+                               np.around(X, decimals=1))
 
 
 def test_kw_arg_reset():
@@ -115,4 +122,4 @@ def test_kw_arg_reset():
 
     # Test that rounding is correct
     testing.assert_array_equal(F.transform(X),
-                                  np.around(X, decimals=1))
+                               np.around(X, decimals=1))
