added y argument to fit methods

added optional y argument to fit methods of TransformerPipeline and
DataFrameMapper

Author: vzaretsk

sklearn_pandas/dataframe_mapper.py

@@ -130,16 +130,27 @@ def fit(self, X, y=None):
         Fit a transformation from the pipeline
 
         X       the data to fit
+
+        y       the target vector relative to X, optional
+
         """
         for columns, transformers in self.features:
             if transformers is not None:
-                transformers.fit(self._get_col_subset(X, columns))
+                try:
+                    transformers.fit(self._get_col_subset(X, columns), y)
+                except TypeError:
+                    # fit takes only one argument
+                    transformers.fit(self._get_col_subset(X, columns))
 
         # handle features not explicitly selected
         if self.default:  # not False and not None
-            self.default.fit(
-                self._get_col_subset(X, self._unselected_columns(X))
-            )
+            try:
+                self.default.fit(
+                    self._get_col_subset(X, self._unselected_columns(X)), y)
+            except TypeError:
+                # fit takes only one argument
+                self.default.fit(
+                    self._get_col_subset(X, self._unselected_columns(X)))
         return self
 
     def transform(self, X):

sklearn_pandas/pipeline.py

@@ -5,10 +5,11 @@
 
 class TransformerPipeline(Pipeline):
     """
-    Pipeline that expects all steps to be transformers taking a single argument
+    Pipeline that expects all steps to be transformers taking a single X argument,
+    an optional y argument,
     and having fit and transform methods.
 
-    Code is copied from sklearn's Pipeline, leaving out the `y=None` argument.
+    Code is copied from sklearn's Pipeline
     """
     def __init__(self, steps):
         names, estimators = zip(*steps)
@@ -31,31 +32,50 @@ def __init__(self, steps):
                             "'%s' (type %s) doesn't)"
                             % (estimator, type(estimator)))
 
-    def _pre_transform(self, X, **fit_params):
+    def _pre_transform(self, X, y=None, **fit_params):
         fit_params_steps = dict((step, {}) for step, _ in self.steps)
         for pname, pval in six.iteritems(fit_params):
             step, param = pname.split('__', 1)
             fit_params_steps[step][param] = pval
         Xt = X
         for name, transform in self.steps[:-1]:
             if hasattr(transform, "fit_transform"):
-                Xt = transform.fit_transform(Xt, **fit_params_steps[name])
+                try:
+                    Xt = transform.fit_transform(Xt, y, **fit_params_steps[name])
+                except TypeError:
+                    # fit takes only one argument
+                    Xt = transform.fit_transform(Xt, **fit_params_steps[name])
             else:
-                Xt = transform.fit(Xt, **fit_params_steps[name]) \
-                              .transform(Xt)
+                try:
+                    Xt = transform.fit(Xt, y, **fit_params_steps[name]).transform(Xt)
+                except TypeError:
+                    # fit takes only one argument
+                    Xt = transform.fit(Xt, **fit_params_steps[name]).transform(Xt)
         return Xt, fit_params_steps[self.steps[-1][0]]
 
-    def fit(self, X, **fit_params):
-        Xt, fit_params = self._pre_transform(X, **fit_params)
-        self.steps[-1][-1].fit(Xt, **fit_params)
+    def fit(self, X, y=None, **fit_params):
+        Xt, fit_params = self._pre_transform(X, y, **fit_params)
+        try:
+            self.steps[-1][-1].fit(Xt, y, **fit_params)
+        except TypeError:
+            # fit takes only one argument
+            self.steps[-1][-1].fit(Xt, **fit_params)
         return self
 
-    def fit_transform(self, X, **fit_params):
-        Xt, fit_params = self._pre_transform(X, **fit_params)
+    def fit_transform(self, X, y=None, **fit_params):
+        Xt, fit_params = self._pre_transform(X, y, **fit_params)
         if hasattr(self.steps[-1][-1], 'fit_transform'):
-            return self.steps[-1][-1].fit_transform(Xt, **fit_params)
+            try:
+                return self.steps[-1][-1].fit_transform(Xt, y, **fit_params)
+            except TypeError:
+                # fit_transform takes only one argument
+                return self.steps[-1][-1].fit_transform(Xt, **fit_params)
         else:
-            return self.steps[-1][-1].fit(Xt, **fit_params).transform(Xt)
+            try:
+                return self.steps[-1][-1].fit(Xt, y, **fit_params).transform(Xt)
+            except:
+                # fit takes only one argument
+                return self.steps[-1][-1].fit(Xt, **fit_params).transform(Xt)
 
 
 def make_transformer_pipeline(*steps):