diff --git a/src/elexmodel/handlers/data/Estimandizer.py b/src/elexmodel/handlers/data/Estimandizer.py index f45832b4..b86c5f33 100644 --- a/src/elexmodel/handlers/data/Estimandizer.py +++ b/src/elexmodel/handlers/data/Estimandizer.py @@ -19,11 +19,14 @@ def add_estimand_results(self, data_df, estimands, historical): for estimand in estimands: results_col = f"{RESULTS_PREFIX}{estimand}" if results_col not in data_df.columns: - if historical and f"{BASELINE_PREFIX}{estimand}" in data_df.columns: - data_df[results_col] = nan - else: - # will raise a KeyError if a function with the same name as `estimand` doesn't exist + # will raise a KeyError if a function with the same name as `estimand` doesn't exist + try: data_df = globals()[estimand](data_df, RESULTS_PREFIX) + except KeyError as e: + if historical and f"{BASELINE_PREFIX}{estimand}" in data_df.columns: + data_df[results_col] = nan + else: + raise e columns_to_return.append(results_col) results_column_names = [x for x in data_df.columns if x.startswith(RESULTS_PREFIX)]