fix test

ihmeuw · Nov 1, 2024 · 33e45a0 · 33e45a0
1 parent c3ad123
commit 33e45a0
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/src/pseudopeople/constants/data_values.py b/src/pseudopeople/constants/data_values.py
@@ -17,6 +17,7 @@
     "White": -0.0188,
 }
 
+# mypy cannot currently type pd.Interval (11/01/24)
 DO_NOT_RESPOND_AGE_INTERVALS: list[pd.Interval] = [  # type: ignore [type-arg]
     # Intervals should include their lower bound
     pd.Interval(0, 5, closed="left"),

diff --git a/src/pseudopeople/interface.py b/src/pseudopeople/interface.py
@@ -2,7 +2,7 @@
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Iterable, Literal, cast
+from typing import Any, Literal, cast
 
 import pandas as pd
 from loguru import logger

diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -258,9 +258,8 @@ def test_two_d_array_choice(
     # Assert proportions
     for sport in best_sports:
         num_choices = 3 if sport != "Baseball" else 2
-        explicit_idx = options.index == sport
-        cols_to_include = options.loc[sport].notna().index.tolist()
-        teams = options.loc[explicit_idx, cols_to_include].values.flatten()
+        # drop NaNs from possible teams
+        teams = [col for col in options.loc[sport] if isinstance(col, str)]
         for team in teams:
             fuzzy_checker.fuzzy_assert_proportion(
                 name="test_two_d_array_choice",