From 6f1c79ef887f09a700d91cc58365ae1672c54970 Mon Sep 17 00:00:00 2001 From: Nicolas Schmid Date: Wed, 16 Aug 2023 12:58:09 +0200 Subject: [PATCH] fix: catalog validation and inplace for strip --- catalog_tools/catalog.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/catalog_tools/catalog.py b/catalog_tools/catalog.py index 765dee2..d1fce18 100644 --- a/catalog_tools/catalog.py +++ b/catalog_tools/catalog.py @@ -7,7 +7,7 @@ def _check_required_cols(df, required_cols=REQUIRED_COLS): - if not set(REQUIRED_COLS).issubset(set(df.columns)): + if not set(required_cols).issubset(set(df.columns)): return False return True @@ -58,5 +58,9 @@ def _constructor(self): @require_cols def strip(self, inplace=False): - self.drop(columns=set(self.columns).difference(set(REQUIRED_COLS)), - inplace=inplace) + """Remove all columns except the required ones. + """ + df = self.drop(columns=set(self.columns).difference(set(REQUIRED_COLS)), + inplace=inplace) + if not inplace: + return df