refactored to use _call_fit

refactored pipeline and dataframe_mapper to use _call_fit helper
function, consolidating code to handle fit with 1 or 2 arguments

Author: vzaretsk

sklearn_pandas/dataframe_mapper.py

@@ -5,7 +5,7 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 
 from .cross_validation import DataWrapper
-from .pipeline import make_transformer_pipeline
+from .pipeline import make_transformer_pipeline, _call_fit
 
 # load in the correct stringtype: str for py3, basestring for py2
 string_types = str if sys.version_info >= (3, 0) else basestring
@@ -136,21 +136,13 @@ def fit(self, X, y=None):
         """
         for columns, transformers in self.features:
             if transformers is not None:
-                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))
+                _call_fit(transformers.fit,
+                          self._get_col_subset(X, columns), y)
 
         # handle features not explicitly selected
         if self.default:  # not False and not None
-            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)))
+            _call_fit(self.default.fit,
+                      self._get_col_subset(X, self._unselected_columns(X)), y)
         return self
 
     def transform(self, X):

sklearn_pandas/pipeline.py

@@ -3,6 +3,30 @@
 from sklearn.utils import tosequence
 
 
+def _call_fit(fit_method, X, y=None, **kwargs):
+    """
+    helper function, calls the fit or fit_transform method with the correct
+    number of parameters
+
+    fit_method: fit or fit_transform method of the transformer
+    X: the data to fit
+    y: the target vector relative to X, optional
+    kwargs: any keyword arguments to the fit method
+
+    return: the result of the fit or fit_transform method
+
+    WARNING: if this function raises a TypeError exception, test the fit
+    or fit_transform method passed to it in isolation as _call_fit will not
+    distinguish TypeError due to incorrect number of arguments from
+    other TypeError
+    """
+    try:
+        return fit_method(X, y, **kwargs)
+    except TypeError:
+        # fit takes only one argument
+        return fit_method(X, **kwargs)
+
+
 class TransformerPipeline(Pipeline):
     """
     Pipeline that expects all steps to be transformers taking a single X argument,
@@ -40,42 +64,26 @@ def _pre_transform(self, X, y=None, **fit_params):
         Xt = X
         for name, transform in self.steps[:-1]:
             if hasattr(transform, "fit_transform"):
-                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])
+                Xt = _call_fit(transform.fit_transform,
+                               Xt, y, **fit_params_steps[name])
             else:
-                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)
+                Xt = _call_fit(transform.fit,
+                               Xt, y, **fit_params_steps[name]).transform(Xt)
         return Xt, fit_params_steps[self.steps[-1][0]]
 
     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)
+        _call_fit(self.steps[-1][-1].fit, Xt, y, **fit_params)
         return self
 
     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'):
-            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)
+            return _call_fit(self.steps[-1][-1].fit_transform,
+                             Xt, y, **fit_params)
         else:
-            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)
+            return _call_fit(self.steps[-1][-1].fit,
+                             Xt, y, **fit_params).transform(Xt)
 
 
 def make_transformer_pipeline(*steps):