pyright: 15 errors to 0, by narrowing where the code already guarantees it

ethnicolr/dict_models.py

@@ -24,6 +24,7 @@
 import logging
 import sys
 from statistics import NormalDist
+from typing import cast
 
 import numpy as np
 import pandas as pd
@@ -100,14 +101,17 @@ def census_marginal(cls) -> np.ndarray:
             )
             marginal = weights.sum(axis=0)
             cls._census_marginal = marginal / marginal.sum()
+        # Set by the branch above on first call, cached after.
+        assert cls._census_marginal is not None
         return cls._census_marginal
 
     @classmethod
     def rosenman(cls, which: str) -> pd.DataFrame:
         if which not in cls._rosenman:
             path = ROSENMAN_FIRST if which == "first" else ROSENMAN_LAST
             df = pd.read_csv(path).dropna(subset=["name"])
-            cls._rosenman[which] = df.set_index("name")[VOTER_CATS]
+            # A list key always yields a DataFrame; the stubs widen it.
+            cls._rosenman[which] = cast(pd.DataFrame, df.set_index("name")[VOTER_CATS])
         return cls._rosenman[which]
 
     @classmethod
@@ -164,7 +168,7 @@ def census_fn(
     df = EthnicolrModelClass.test_and_norm_df(df, fname_col)
 
     table = _Tables.census_first()
-    keys = _norm_names(df[fname_col])
+    keys = _norm_names(cast(pd.Series, df[fname_col]))
     matched = table.reindex(keys)
 
     rdf = df.copy()
@@ -181,7 +185,7 @@ def census_fn(
             rdf[f"{col}_lb"] = (lb * 100).round(2)
             rdf[f"{col}_ub"] = (ub * 100).round(2)
 
-    matched_n = int(matched[CENSUS_PCT_COLS[0]].notna().sum())
+    matched_n = int(cast(pd.Series, matched[CENSUS_PCT_COLS[0]]).notna().sum())
     logger.info(f"Matched {matched_n} of {len(rdf)} first names")
     return rdf
 
@@ -235,8 +239,8 @@ def pred_census_name(
         raise ValueError("lname_col and fname_col must exist in the DataFrame")
 
     rdf = df.copy()
-    last_keys = _norm_names(rdf[lname_col])
-    first_keys = _norm_names(rdf[fname_col])
+    last_keys = _norm_names(cast(pd.Series, rdf[lname_col]))
+    first_keys = _norm_names(cast(pd.Series, rdf[fname_col]))
 
     last_table = _Tables.census_last(year)
     first_table = _Tables.census_first()
@@ -354,8 +358,8 @@ def pred_voter_name(
         raise ValueError("lname_col and fname_col must exist in the DataFrame")
 
     rdf = df.copy()
-    last_keys = _norm_names(rdf[lname_col])
-    first_keys = _norm_names(rdf[fname_col])
+    last_keys = _norm_names(cast(pd.Series, rdf[lname_col]))
+    first_keys = _norm_names(cast(pd.Series, rdf[fname_col]))
 
     last_table = _Tables.rosenman("last")
     first_table = _Tables.rosenman("first")

ethnicolr/pred_census_ln.py

@@ -213,11 +213,15 @@ def pred_census_ln(
             "this model has no calibration stats file; run "
             "scripts/model-training/calibrate_model.py"
         )
-    if coverage is not None and f"{coverage:.2f}" not in stats["conformal_quantiles"]:
-        raise ValueError(
-            f"coverage must be one of {sorted(stats['conformal_quantiles'])}, "
-            f"got {coverage}"
-        )
+    if coverage is not None:
+        # Guaranteed by the check above: stats is None only when neither prior
+        # nor coverage was given.
+        assert stats is not None
+        if f"{coverage:.2f}" not in stats["conformal_quantiles"]:
+            raise ValueError(
+                f"coverage must be one of {sorted(stats['conformal_quantiles'])}, "
+                f"got {coverage}"
+            )
 
     logger.info(f"Predicting {len(df)} names using Census {year} PyTorch model")
 
@@ -255,6 +259,7 @@ def pred_census_ln(
             logits = model(X_tensor)
             mean_probs = torch.softmax(logits / temperature, dim=1).cpu().numpy()
             if prior is not None:
+                assert stats is not None  # guaranteed by the check above
                 mean_probs = apply_prior(
                     mean_probs, RACES, prior, stats["train_class_distribution"]
                 )
@@ -291,6 +296,7 @@ def pred_census_ln(
     result["race"] = [RACES[i] for i in pred_indices]
 
     if coverage is not None:
+        assert stats is not None  # guaranteed by the check above
         qhat = stats["conformal_quantiles"][f"{coverage:.2f}"]
         result["race_set"] = conformal_sets(mean_probs, RACES, qhat)
 

ethnicolr/pred_wiki_origin.py

@@ -73,6 +73,12 @@ def pred_wiki_origin(
             + working[fname_col].fillna("").astype(str).str.strip()
         ).str.strip()
 
+        # The base declares VOCABFN/RACEFN as `str | None` because non-LSTM
+        # models have neither. This is an LSTM model and sets both; narrowing
+        # them in the subclass does not work, since a mutable ClassVar is
+        # invariant and `str` is not assignable to `str | None`.
+        assert cls.VOCABFN is not None and cls.RACEFN is not None
+
         rdf = cls.transform_and_pred(
             df=working,
             newnamecol="__name",