Closes #100: exclude_columns not working for ColumnsBasedPipelineStage

pdpipe · May 12, 2022 · a382412 · a382412
1 parent b0f84bc
commit a382412
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/pdpipe/core.py b/pdpipe/core.py
@@ -599,8 +599,8 @@ def __init__(
             none_columns='error', **kwargs):
         self._exclude_columns = exclude_columns
         if exclude_columns:
-            self._exclude_columns = self._interpret_columns_param(
-                exclude_columns)
+            self._exclude_columns, self._exc_col_str = \
+                self._interpret_columns_param(exclude_columns)
         self._none_error = False
         self._none_cols = None
         # handle none_columns
@@ -625,13 +625,14 @@ def __init__(
         # done handling none_columns
         self._col_arg, self._col_str = self._interpret_columns_param(
             columns, self._none_error, none_columns=self._none_cols)
+        self._final_col_str = '{self._col_str} (except {self._exc_col_str})'
         if (kwargs.get('desc') is None) and desc_temp:
-            kwargs['desc'] = desc_temp.format(self._col_str)
+            kwargs['desc'] = desc_temp.format(self._final_col_str)
         if kwargs.get('exmsg') is None:
             kwargs['exmsg'] = (
                 'Pipeline stage failed because not all columns {} '
                 'were found in the input dataframe.'
-            ).format(self._col_str)
+            ).format(self._final_col_str)
         super().__init__(**kwargs)
 
     def _is_fittable(self):