Works but it's not really addressing the problem

washingtonpost · Sep 18, 2023 · 01d1b2a · 01d1b2a
1 parent b96ba77
commit 01d1b2a
Showing 1 changed file with 7 additions and 4 deletions.
diff --git 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)]