Update wrapper

ihmeuw · Nov 7, 2024 · 793ac11 · 793ac11
1 parent e27aee7
commit 793ac11
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/vivarium/framework/results/observation.py b/src/vivarium/framework/results/observation.py
@@ -139,16 +139,21 @@ def __init__(
         to_observe: Callable[[Event], bool] = lambda event: True,
     ):
         def _wrap_results_gatherer(
-            df: pd.DataFrame, _: tuple[str, ...] | None
+            df: pd.DataFrame | DataFrameGroupBy[str], _: tuple[str, ...] | None
         ) -> pd.DataFrame:
+            if isinstance(df, DataFrameGroupBy):
+                raise TypeError(
+                    "Must provide a dataframe to an UnstratifiedObservation. "
+                    f"Provided DataFrameGroupBy instead."
+                )
             return results_gatherer(df)
 
         super().__init__(
             name=name,
             pop_filter=pop_filter,
             when=when,
             results_initializer=self.create_empty_df,
-            results_gatherer=_wrap_results_gatherer,  # type: ignore [arg-type]
+            results_gatherer=_wrap_results_gatherer,
             results_updater=results_updater,
             results_formatter=results_formatter,
             stratifications=None,