From 01d1b2ab47f7d60f90c7cb90816943b7f90e431b Mon Sep 17 00:00:00 2001 From: Diane Napolitano Date: Mon, 18 Sep 2023 14:47:21 -0400 Subject: [PATCH] Works but it's not really addressing the problem --- src/elexmodel/handlers/data/Estimandizer.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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)]