diff --git a/financedatabase/helpers.py b/financedatabase/helpers.py index e6ee2d265..99e5e92d8 100644 --- 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