fix mypy error

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

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/pyfixest/estimation/feols_.py b/pyfixest/estimation/feols_.py
@@ -1844,9 +1844,7 @@ def predict(
             fixef_dicts = {}
             for f in fvals:
                 fdict = self._fixef_dict[f"C({f})"]
-                omitted_cat = {
-                    str(x) for x in self._data[f].unique() if str(x) not in fdict
-                }
+                omitted_cat = {str(x) for x in self._data[f].unique() if x not in fdict}
                 if omitted_cat:
                     fdict.update({x: 0 for x in omitted_cat})
                 fixef_dicts[f"C({f})"] = fdict

diff --git a/pyfixest/utils/dev_utils.py b/pyfixest/utils/dev_utils.py
@@ -167,7 +167,7 @@ def _drop_cols(_data: pd.DataFrame, na_index: np.ndarray):
         return _data
 
 
-def _extract_variable_level(fe_string: str):
+def _extract_variable_level(fe_string: str) -> tuple[str, str]:
     """
     Extract the variable and level from a given string.