Update type stubs

Author: tmke8

ethicml/implementations/agarwal.py

@@ -26,7 +26,7 @@
 )
 
 if TYPE_CHECKING:
-    from fairlearn.reductions import ExponentiatedGradient
+    from fairlearn.reductions import ExponentiatedGradient  # pyright: ignore
 
     from ethicml.models.inprocess.agarwal_reductions import AgarwalArgs
     from ethicml.models.inprocess.in_subprocess import InAlgoArgs
@@ -36,7 +36,7 @@
 def fit(train: DataTuple, args: AgarwalArgs, seed: int = 888) -> ExponentiatedGradient:
     """Fit a model."""
     try:
-        from fairlearn.reductions import (
+        from fairlearn.reductions import (  # pyright: ignore
             DemographicParity,
             EqualizedOdds,
             ExponentiatedGradient,
@@ -81,7 +81,7 @@ def fit(train: DataTuple, args: AgarwalArgs, seed: int = 888) -> ExponentiatedGr
     exponentiated_gradient.fit(data_x, data_y, sensitive_features=data_a)
 
     min_class_label = train.y.min()
-    exponentiated_gradient.min_class_label = min_class_label
+    exponentiated_gradient.min_class_label = min_class_label  # pyright: ignore
 
     return exponentiated_gradient
 
@@ -92,7 +92,7 @@ def predict(exponentiated_gradient: ExponentiatedGradient, test: TestTuple) -> p
     preds = pd.DataFrame(randomized_predictions, columns=["preds"])
 
     if (min_val := preds["preds"].min()) != preds["preds"].max():
-        preds = preds.replace(min_val, exponentiated_gradient.min_class_label)
+        preds = preds.replace(min_val, exponentiated_gradient.min_class_label)  # pyright: ignore
     return preds
 
 
@@ -120,7 +120,7 @@ def main() -> None:
     in_algo_args: InAlgoArgs = json.loads(sys.argv[1])
     flags: AgarwalArgs = json.loads(sys.argv[2])
     try:
-        import cloudpickle
+        import cloudpickle  # pyright: ignore
 
         # Need to install cloudpickle for now. See https://github.com/fairlearn/fairlearn/issues/569
     except ImportError as e:

ethicml/metrics/accuracy.py

@@ -30,23 +30,23 @@ def score(self, prediction: Prediction, actual: EvalTuple) -> float:
 class Accuracy(SklearnMetric):
     """Classification accuracy."""
 
-    sklearn_metric = (accuracy_score,)
+    sklearn_metric = (accuracy_score,)  # type: ignore[assignment]
     _name: ClassVar[str] = "Accuracy"
 
 
 @dataclass
 class F1(SklearnMetric):
     """F1 score: harmonic mean of precision and recall."""
 
-    sklearn_metric = (f1_score,)
+    sklearn_metric = (f1_score,)  # type: ignore[assignment]
     _name: ClassVar[str] = "F1"
 
 
 @dataclass
 class RobustAccuracy(SklearnMetric):
     """Minimum Classification accuracy across S-groups."""
 
-    sklearn_metric = (accuracy_score,)
+    sklearn_metric = (accuracy_score,)  # type: ignore[assignment]
     apply_per_sensitive: ClassVar[bool] = False
     _name: ClassVar[str] = "Robust Accuracy"
 

ethicml/metrics/dependence_measures.py

@@ -49,7 +49,7 @@ class NMI(_DependenceMeasure):
     @override
     def score(self, prediction: Prediction, actual: EvalTuple) -> float:
         base_values = actual.y if self.base is DependencyTarget.y else actual.s
-        return normalized_mutual_info_score(
+        return normalized_mutual_info_score(  # type: ignore[return-value]
             base_values.to_numpy().ravel(),
             prediction.hard.to_numpy().ravel(),
             average_method="arithmetic",

ethicml/models/inprocess/kamiran.py

@@ -130,7 +130,7 @@ def _predict(self, model: LinearModel, test: TestTuple) -> Prediction:
                 hard=pd.Series(model.predict(test.x.to_numpy())), info=self.hyperparameters
             )
         return SoftPrediction(
-            soft=model.predict_proba(test.x.to_numpy()),  # type: ignore[arg-type]
+            soft=model.predict_proba(test.x.to_numpy()),
             info=self.hyperparameters,
         )
 

ethicml/models/inprocess/logistic_regression.py

@@ -55,7 +55,7 @@ def run(self, train: DataTuple, test: TestTuple, seed: int = 888) -> SoftPredict
         )
         clf.fit(train.x, train.y.to_numpy().ravel())
         return SoftPrediction(
-            soft=clf.predict_proba(test.x.to_numpy()),  # type: ignore[arg-type]
+            soft=clf.predict_proba(test.x.to_numpy()),
             info=self.hyperparameters,
         )
 
@@ -88,9 +88,9 @@ def fit(self, train: DataTuple, seed: int = 888) -> LRCV:
     @override
     def predict(self, test: TestTuple) -> Prediction:
         params = self.hyperparameters
-        params["C"] = self.clf.C_[0]  # type: ignore[attr-defined]
+        params["C"] = self.clf.C_[0]
         return SoftPrediction(
-            soft=self.clf.predict_proba(test.x.to_numpy()),  # type: ignore[arg-type]
+            soft=self.clf.predict_proba(test.x.to_numpy()),
             info=params,
         )
 
@@ -103,8 +103,8 @@ def run(self, train: DataTuple, test: TestTuple, seed: int = 888) -> Prediction:
         )
         clf.fit(train.x, train.y.to_numpy().ravel())
         params = self.hyperparameters
-        params["C"] = clf.C_[0]  # type: ignore[attr-defined]
+        params["C"] = clf.C_[0]
         return SoftPrediction(
-            soft=clf.predict_proba(test.x.to_numpy()),  # type: ignore[arg-type]
+            soft=clf.predict_proba(test.x.to_numpy()),
             info=params,
         )

ethicml/plot/plotting.py

@@ -32,9 +32,7 @@ def _maybe_tsne(data: DataTuple) -> tuple[pd.DataFrame, str, str]:
     if len(columns) > 2:
         from sklearn.manifold import TSNE
 
-        tsne_embeddings = TSNE(n_components=2, random_state=0).fit_transform(
-            data.x,  # type: ignore[arg-type]
-        )
+        tsne_embeddings = TSNE(n_components=2, random_state=0).fit_transform(data.x)
         amalgamated = pd.concat(
             [pd.DataFrame(tsne_embeddings, columns=["tsne1", "tsne2"]), data.s, data.y],
             axis="columns",

ethicml/preprocessing/splits.py

@@ -116,8 +116,8 @@ def train_test_split(
     pd.testing.assert_index_equal(train.x.columns, x_columns)
     pd.testing.assert_index_equal(test.x.columns, x_columns)
 
-    assert isinstance(train.s, pd.Series)
-    assert isinstance(test.s, pd.Series)
+    assert isinstance(train.s, pd.Series)  # type: ignore[unreachable]
+    assert isinstance(test.s, pd.Series)  # type: ignore[unreachable]
     assert train.s.name == data.s_column
     assert test.s.name == data.s_column
 

poetry.lock

@@ -54,7 +54,7 @@ mypy = ">=0.990"
 pre-commit = "^2.20.0"
 pytest = ">=6.2.2,<8.0.0"
 pytest-cov = ">=2.6,<4.0"
-python-type-stubs = {git = "https://github.com/wearepal/python-type-stubs.git", rev = "417b03b"}
+python-type-stubs = {git = "https://github.com/wearepal/python-type-stubs.git", rev = "2ea8053"}
 pandas-stubs = ">=1.4.2.220626"
 omegaconf = ">=2.2.1"
 pytest-xdist = "^2.5.0"
@@ -173,7 +173,6 @@ module = [
     "pytest.*",
     "setuptools.*",
     "scipy.spatial.distance",
-    "sklearn.*",
 ]
 ignore_missing_imports = true