fix: catalog validation and inplace for strip

swiss-seismological-service · Aug 16, 2023 · 6f1c79e · 6f1c79e
1 parent b4869b3
commit 6f1c79e
Showing 1 changed file with 7 additions and 3 deletions.
diff --git 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