fix mypy error

py-econometrics · Dec 25, 2024 · c81f021 · c81f021
1 parent e501b54
commit c81f021
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/pyfixest/estimation/feols_.py b/pyfixest/estimation/feols_.py
@@ -3,7 +3,7 @@
 import re
 import warnings
 from importlib import import_module
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import numba as nb
 import numpy as np
@@ -1844,9 +1844,9 @@ def predict(
             fixef_dicts = {}
             for f in fvals:
                 fdict = self._fixef_dict[f"C({f})"]
-                omitted_cat = set(self._data[f].unique().astype(str).tolist()) - set(
-                    fdict.keys()
-                )
+                omitted_cat = {
+                    str(x) for x in self._data[f].unique() if str(x) not in fdict
+                }
                 if omitted_cat:
                     fdict.update({x: 0 for x in omitted_cat})
                 fixef_dicts[f"C({f})"] = fdict
@@ -2514,7 +2514,7 @@ def _get_vcov_type(vcov: str, fval: str):
 
 def _drop_multicollinear_variables(
     X: np.ndarray, names: list[str], collin_tol: float
-) -> tuple[np.ndarray, list[str], list[str], list[int]]:
+) -> Any:
     """
     Check for multicollinearity in the design matrices X and Z.
 

diff --git a/pyfixest/estimation/model_matrix_fixest_.py b/pyfixest/estimation/model_matrix_fixest_.py
@@ -209,14 +209,15 @@ def model_matrix_fixest(
     }
 
 
-def _get_na_index(N: int, Y_index: pd.Series) -> np.ndarray:
+def _get_na_index(N: int, Y_index: pd.Index) -> np.ndarray:
     all_indices = np.arange(N)
 
     max_index = all_indices.max() + 1
     mask = np.ones(max_index, dtype=bool)
     Y_index_arr = Y_index.to_numpy()
+
     mask[Y_index_arr] = False
-    na_index = np.nonzero(mask)[0]
+    na_index = np.flatnonzero(mask)
     return na_index