From 8beb8e386255bdd1ed5d9783d7414a2384cc3d17 Mon Sep 17 00:00:00 2001 From: "Hamshaw, Scott Douglas" Date: Wed, 10 Apr 2024 13:43:57 -0400 Subject: [PATCH] update filtering of approved data to search for any match --- hyswap/utils.py | 4 +--- tests/test_utils.py | 5 +++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/hyswap/utils.py b/hyswap/utils.py index 548243c..32a314e 100644 --- a/hyswap/utils.py +++ b/hyswap/utils.py @@ -41,9 +41,7 @@ def filter_approved_data(data, filter_column=None): """ if filter_column is None: raise ValueError("filter_column must be specified.") - return data.loc[((data[filter_column] == "A") | - (data[filter_column] == "A, e") | - (data[filter_column] == "A, R"))] + return data[data[filter_column].str.contains("A")] def rolling_average(df, data_column_name, data_type, diff --git a/tests/test_utils.py b/tests/test_utils.py index 2bd41e9..d66efa7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -8,10 +8,11 @@ class TestDataFilters: def test_filter_approved_data(self): """Test the filter_approved_data function.""" - data = {"a": ["A", "A", "P", "P"], "b": [1, 2, 3, 4]} + data = {"a": ["A", "A, e", "A, R", "A, [4]", "P", "P"], + "b": [1, 2, 3, 4, 5, 6]} data_df = pd.DataFrame(data) df = utils.filter_approved_data(data_df, filter_column="a") - assert df["b"].tolist() == [1, 2] + assert df["b"].tolist() == [1, 2, 3, 4] def test_filter_approved_data_error(self): """Test the filter_approved_data function."""