From 3d84dff865537b40f1d5aaf4bd4169098a37ef71 Mon Sep 17 00:00:00 2001
From: Alexander Fischer <alexander-fischer1801@t-online.de>
Date: Wed, 25 Dec 2024 13:47:42 +0100
Subject: [PATCH] fix mypy error

---
 pyfixest/estimation/feols_.py | 4 +---
 pyfixest/utils/dev_utils.py   | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/pyfixest/estimation/feols_.py b/pyfixest/estimation/feols_.py
index f778a686a..e0405c32a 100644
--- 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
index bc8966f04..5f200624a 100644
--- 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.