expression plugin can now drop indexes as well

CountESS-Project · Jun 27, 2023 · 308da5a · 308da5a
1 parent f310886
commit 308da5a
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/countess/plugins/expression.py b/countess/plugins/expression.py
@@ -46,9 +46,16 @@ def run_df(self, df, logger: Logger) -> pd.DataFrame:
 
         df = process(df, codes, logger)
 
-        drop_columns = [
+        drop_names = [
             col for col, param in zip(self.input_columns, self.parameters["drop"]) if param.value
         ]
 
-        print(drop_columns)
-        return df.drop(columns=drop_columns)
+        drop_indexes = [ col for col in drop_names if col in df.index.names ]
+        if drop_indexes:
+            df = df.reset_index(drop_indexes, drop=True)
+
+        drop_columns = [ col for col in drop_names if col in df.columns ]
+        if drop_columns:
+            df = df.drop(columns=drop_columns)
+
+        return df