From 43b2fa04b0c1b673c240b224c09ef705b2c566f9 Mon Sep 17 00:00:00 2001 From: R-Palazzo Date: Thu, 24 Aug 2023 15:31:20 +0200 Subject: [PATCH] drop _store_columns change --- rdt/transformers/base.py | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/rdt/transformers/base.py b/rdt/transformers/base.py index 5071358c5..c9128d340 100644 --- a/rdt/transformers/base.py +++ b/rdt/transformers/base.py @@ -259,25 +259,16 @@ def get_output_columns(self): return list(self._get_output_to_property('sdtype')) def _store_columns(self, columns, data): - """Store the column names to be transformed. + if isinstance(columns, tuple) and columns not in data: + columns = list(columns) + elif not isinstance(columns, list): + columns = [columns] - Args: - columns (str, list, tuple): - Column name(s) to be transformed. - data (pandas.DataFrame): - The entire table. - """ - column_names = columns - if isinstance(column_names, tuple) and column_names not in data: - column_names = list(column_names) - elif not isinstance(column_names, list): - column_names = [column_names] - - missing = set(column_names) - set(data.columns) + missing = set(columns) - set(data.columns) if missing: raise KeyError(f'Columns {missing} were not present in the data.') - self.columns = column_names + self.columns = columns @staticmethod def _get_columns_data(data, columns):