Add in the option to search for multiple values instead of just one

JerBouma · Aug 24, 2023 · e4f4178 · e4f4178
1 parent f93f55a
commit e4f4178
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/financedatabase/helpers.py b/financedatabase/helpers.py
@@ -84,15 +84,25 @@ def search(self, **kwargs: str) -> pd.DataFrame:
                     ]
             elif key == "index":
                 # Look into the index of the DataFrame and search accordingly
-                data_filter = data_filter[
-                    data_filter.index.str.contains(value, na=False)
-                ]
+                if isinstance(value, (list, pd.Index)):
+                    data_filter = data_filter[
+                        data_filter.index.isin(value)
+                    ]
+                else:
+                    data_filter = data_filter[
+                        data_filter.index.str.contains(value, na=False)
+                    ]
             elif key not in data_filter.columns:
                 print(f"{key} is not a valid column.")
             else:
-                data_filter = data_filter[
-                    data_filter[key].str.contains(value, case=case_sensitive, na=False)
-                ]
+                if isinstance(value, list):
+                    data_filter = data_filter[
+                        data_filter[key].isin(value)
+                    ]
+                else:
+                    data_filter = data_filter[
+                        data_filter[key].str.contains(value, case=case_sensitive, na=False)
+                    ]
 
         return data_filter