Fix typing

Author: dustalov

.pre-commit-config.yaml

@@ -1,25 +1,25 @@
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v5.0.0
+  rev: v6.0.0
   hooks:
   - id: trailing-whitespace
   - id: end-of-file-fixer
   - id: check-yaml
   - id: check-added-large-files
 - repo: https://github.com/psf/black
-  rev: 24.10.0
+  rev: 25.9.0
   hooks:
   - id: black
 - repo: https://github.com/pycqa/isort
-  rev: 5.13.2
+  rev: 7.0.0
   hooks:
   - id: isort
 - repo: https://github.com/pycqa/flake8
-  rev: 7.1.1
+  rev: 7.3.0
   hooks:
   - id: flake8
 - repo: https://github.com/asottile/pyupgrade
-  rev: v3.19.1
+  rev: v3.21.0
   hooks:
   - id: pyupgrade
 - repo: https://github.com/nbQA-dev/nbQA

crowdkit/aggregation/classification/dawid_skene.py

@@ -1,6 +1,6 @@
 __all__ = ["DawidSkene", "OneCoinDawidSkene"]
 
-from typing import Any, List, Literal, Optional, cast
+from typing import Any, List, Literal, Optional, Tuple, cast
 
 import attr
 import numpy as np
@@ -235,7 +235,7 @@ def _evidence_lower_bound(
         priors = priors.rename(index={True: "True", False: "False"}, copy=False)
         priors.clip(lower=_EPS, inplace=True)
 
-        joined.loc[:, priors.index] = joined.loc[:, priors.index].add(np.log(priors))  # type: ignore
+        joined.loc[:, priors.index] = joined.loc[:, priors.index].add(np.log(priors))
 
         joined.set_index(["task", "worker"], inplace=True)
         joint_expectation = (
@@ -487,25 +487,29 @@ class OneCoinDawidSkene(DawidSkene):
     """
 
     @staticmethod
-    def _assign_skills(row: "pd.Series[Any]", skills: pd.DataFrame) -> pd.DataFrame:
+    def _assign_skills(
+        row: "pd.Series[Any]", skills: "pd.Series[Any]"
+    ) -> "pd.Series[Any]":
         """
         Assigns user skills to error matrix row by row.
         """
         num_categories = len(row)
         for column_name, _ in row.items():
-            if column_name == row.name[1]:  # type: ignore
-                row[column_name] = skills[row.name[0]]  # type: ignore
+            if column_name == cast(Tuple[Any, Any], row.name)[1]:
+                row[column_name] = skills.loc[cast(Tuple[Any, Any], row.name)[0]]
             else:
-                row[column_name] = (1 - skills[row.name[0]]) / (num_categories - 1)  # type: ignore
-        return row  # type: ignore
+                row[column_name] = (
+                    1 - skills.loc[cast(Tuple[Any, Any], row.name)[0]]
+                ) / (num_categories - 1)
+        return row
 
     @staticmethod
     def _process_skills_to_errors(
         data: pd.DataFrame, probas: pd.DataFrame, skills: "pd.Series[Any]"
     ) -> pd.DataFrame:
         errors = DawidSkene._m_step(data, probas)
 
-        errors = errors.apply(OneCoinDawidSkene._assign_skills, args=(skills,), axis=1)  # type: ignore
+        errors = errors.apply(OneCoinDawidSkene._assign_skills, args=(skills,), axis=1)
         errors.clip(lower=_EPS, upper=1 - _EPS, inplace=True)
 
         return errors

crowdkit/aggregation/classification/glad.py

@@ -226,8 +226,8 @@ def _optimize_df(self, x: npt.NDArray[Any]) -> npt.NDArray[Any]:
         dQalpha, dQbeta = self._gradient_Q(self._current_data)
 
         minus_grad = np.zeros_like(x)
-        minus_grad[: len(self.workers_)] = -dQalpha[self.workers_].values  # type: ignore
-        minus_grad[len(self.workers_) :] = -dQbeta[self.tasks_].values  # type: ignore
+        minus_grad[: len(self.workers_)] = -dQalpha[self.workers_].values  # type: ignore[operator,index]
+        minus_grad[len(self.workers_) :] = -dQbeta[self.tasks_].values  # type: ignore[operator,index]
         return minus_grad
 
     def _update_alphas_betas(
@@ -245,9 +245,9 @@ def _update_alphas_betas(
     def _get_alphas_betas_by_point(
         self, x: npt.NDArray[Any]
     ) -> Tuple["pd.Series[Any]", "pd.Series[Any]"]:
-        alphas = pd.Series(x[: len(self.workers_)], index=self.workers_, name="alpha")  # type: ignore
+        alphas = pd.Series(x[: len(self.workers_)], index=self.workers_, name="alpha")
         alphas.index.name = "worker"
-        betas = pd.Series(x[len(self.workers_) :], index=self.tasks_, name="beta")  # type: ignore
+        betas = pd.Series(x[len(self.workers_) :], index=self.tasks_, name="beta")
         betas.index.name = "task"
         return alphas, betas
 
@@ -268,15 +268,15 @@ def _m_step(self, data: pd.DataFrame) -> pd.DataFrame:
         return self._current_data
 
     def _init(self, data: pd.DataFrame) -> None:
-        self.alphas_ = pd.Series(1.0, index=pd.unique(data.worker))  # type: ignore
-        self.betas_ = pd.Series(1.0, index=pd.unique(data.task))  # type: ignore
+        self.alphas_ = pd.Series(1.0, index=pd.unique(data.worker))
+        self.betas_ = pd.Series(1.0, index=pd.unique(data.task))
         self.tasks_ = pd.unique(data["task"])
         self.workers_ = pd.unique(data["worker"])
         self.priors_ = self.labels_priors
         if self.priors_ is None:
             self.prior_labels_ = pd.unique(data["label"])
             self.priors_ = pd.Series(
-                1.0 / len(self.prior_labels_), index=self.prior_labels_  # type: ignore
+                1.0 / len(self.prior_labels_), index=self.prior_labels_
             )
         else:
             self.prior_labels_ = self.priors_.index  # type: ignore

crowdkit/aggregation/classification/mace.py

@@ -1,6 +1,6 @@
 __all__ = ["MACE"]
 
-from typing import Any, Iterator, List, Optional, Tuple, Union
+from typing import Any, Iterator, List, Optional, Tuple, Union, cast
 
 import attr
 import numpy as np
@@ -24,11 +24,14 @@ def normalize(x: NDArray[np.float64], smoothing: float) -> NDArray[np.float64]:
         np.ndarray: Normalized array
     """
     norm = (x + smoothing).sum(axis=1)
-    return np.divide(
-        x + smoothing,
-        norm[:, np.newaxis],
-        out=np.zeros_like(x),
-        where=~np.isclose(norm[:, np.newaxis], np.zeros_like(norm[:, np.newaxis])),
+    return cast(
+        NDArray[np.float64],
+        np.divide(
+            x + smoothing,
+            norm[:, np.newaxis],
+            out=np.zeros_like(x),
+            where=~np.isclose(norm[:, np.newaxis], np.zeros_like(norm[:, np.newaxis])),
+        ),
     )
 
 
@@ -46,11 +49,14 @@ def variational_normalize(
     """
     norm = (x + hparams).sum(axis=1)
     norm = np.exp(digamma(norm))
-    return np.divide(
-        np.exp(digamma(x + hparams)),
-        norm[:, np.newaxis],
-        out=np.zeros_like(x),
-        where=~np.isclose(norm[:, np.newaxis], np.zeros_like(norm[:, np.newaxis])),
+    return cast(
+        NDArray[np.float64],
+        np.divide(
+            np.exp(digamma(x + hparams)),
+            norm[:, np.newaxis],
+            out=np.zeros_like(x),
+            where=~np.isclose(norm[:, np.newaxis], np.zeros_like(norm[:, np.newaxis])),
+        ),
     )
 
 

crowdkit/aggregation/texts/text_hrrasa.py

@@ -112,16 +112,16 @@ def fit_predict(  # type: ignore
             self._encode_data(data), self._encode_true_objects(true_objects)
         )
         self.texts_ = (
-            hrrasa_results.reset_index()[["task", "output"]]  # type: ignore
+            hrrasa_results.reset_index()[["task", "output"]]
             .rename(columns={"output": "text"})
-            .set_index("task")
+            .set_index("task")["text"]
         )
         return self.texts_
 
     def _encode_data(self, data: pd.DataFrame) -> pd.DataFrame:
         data = data[["task", "worker", "text"]].rename(columns={"text": "output"})
-        data["embedding"] = data.output.apply(self.encoder)  # type: ignore
+        data["embedding"] = data.output.apply(self.encoder)  # type: ignore[arg-type]
         return data
 
     def _encode_true_objects(self, true_objects: "pd.Series[Any]") -> "pd.Series[Any]":
-        return true_objects and true_objects.apply(self.encoder)  # type: ignore
+        return true_objects and true_objects.apply(self.encoder)

crowdkit/aggregation/texts/text_rasa.py

@@ -114,9 +114,9 @@ def fit_predict(  # type: ignore
             self._encode_data(data), self._encode_true_objects(true_objects)
         )
         self.texts_ = (
-            rasa_results.reset_index()[["task", "output"]]  # type: ignore
+            rasa_results.reset_index()[["task", "output"]]
             .rename(columns={"output": "text"})
-            .set_index("task")
+            .set_index("task")["text"]
         )
         return self.texts_
 
@@ -126,4 +126,4 @@ def _encode_data(self, data: pd.DataFrame) -> pd.DataFrame:
         return data
 
     def _encode_true_objects(self, true_objects: "pd.Series[Any]") -> "pd.Series[Any]":
-        return true_objects and true_objects.apply(self.encoder)  # type: ignore
+        return true_objects and true_objects.apply(self.encoder)

crowdkit/datasets/_base.py

@@ -4,11 +4,11 @@
 from os import environ, listdir, makedirs, rename
 from os.path import basename, exists, expanduser, join, splitext
 from shutil import unpack_archive
-from typing import AnyStr, Optional, cast
+from typing import Optional
 from urllib.request import urlretrieve
 
 
-def get_data_dir(data_dir: Optional[AnyStr] = None) -> AnyStr:
+def get_data_dir(data_dir: Optional[str] = None) -> str:
     """Return the path of the crowd-kit data dir.
 
     This folder is used by some large dataset loaders to avoid downloading the
@@ -26,9 +26,7 @@ def get_data_dir(data_dir: Optional[AnyStr] = None) -> AnyStr:
             is `~/crowdkit_data`.
     """
     if data_dir is None:
-        data_dir = cast(
-            AnyStr, environ.get("CROWDKIT_DATA", join("~", "crowdkit_data"))
-        )
+        data_dir = environ.get("CROWDKIT_DATA", join("~", "crowdkit_data"))
         data_dir = expanduser(data_dir)
 
     if not exists(data_dir):

crowdkit/learning/text_summarization.py

@@ -82,7 +82,7 @@ def fit_predict(self, data: pd.DataFrame) -> "pd.Series[Any]":
 
         data = data[["task", "worker", "text"]]
 
-        self.model = self.model.to(self.device)
+        self.model = self.model.to(self.device)  # type: ignore[arg-type]
         self.texts_ = data.groupby("task")["text"].apply(self._aggregate_one)
         return self.texts_
 
@@ -117,5 +117,5 @@ def _generate_output(
         input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to(
             self.device
         )
-        outputs = self.model.generate(input_ids, num_beams=self.num_beams)
-        return cast(str, self.tokenizer.decode(outputs[0], skip_special_tokens=True))
+        outputs = self.model.generate(input_ids, num_beams=self.num_beams)  # type: ignore[operator]
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

pyproject.toml

@@ -96,7 +96,6 @@ version = {attr = "crowdkit.__version__"}
 
 [tool.mypy]
 ignore_missing_imports = true
-plugins = ["numpy.typing.mypy_plugin"]
 strict = true
 
 [tool.isort]

tests/aggregation/test_ds_aggregation.py

@@ -3,7 +3,7 @@
 Testing all boundary conditions and asserts
 """
 
-from typing import Any, List, Literal, Optional, cast
+from typing import Any, List, Literal, Optional
 
 import numpy as np
 import pandas as pd
@@ -461,14 +461,12 @@ def _make_probas(data: List[List[Any]]) -> pd.DataFrame:
     return pd.DataFrame(data, columns=columns).set_index("task")
 
 
-def _make_tasks_labels(data: List[List[Any]]) -> pd.DataFrame:
+def _make_tasks_labels(data: List[List[Any]]) -> "pd.Series[Any]":
     # TODO: should task be indexed?
-    return cast(
-        pd.DataFrame,
+    return (
         pd.DataFrame(data, columns=["task", "label"])
-        .set_index("task")
-        .squeeze()
-        .rename("agg_label"),
+        .set_index("task")["label"]
+        .rename("agg_label")
     )
 
 
@@ -532,7 +530,7 @@ def priors_iter_0() -> "pd.Series[Any]":
 
 
 @pytest.fixture
-def tasks_labels_iter_0() -> pd.DataFrame:
+def tasks_labels_iter_0() -> "pd.Series[Any]":
     return _make_tasks_labels(
         [
             ["t1", "no"],
@@ -581,7 +579,7 @@ def priors_iter_1() -> "pd.Series[Any]":
 
 
 @pytest.fixture
-def tasks_labels_iter_1() -> pd.DataFrame:
+def tasks_labels_iter_1() -> "pd.Series[Any]":
     return _make_tasks_labels(
         [
             ["t1", "yes"],
@@ -670,7 +668,7 @@ def test_dawid_skene_overlap(overlap: int) -> None:
     assert ds.priors_ is not None, "no priors_"
     assert ds.labels_ is not None, "no labels_"
     assert_frame_equal(expected_probas, ds.probas_, check_like=True, atol=0.005)
-    assert_series_equal(expected_labels, ds.labels_, atol=0.005)  # type: ignore
+    assert_series_equal(expected_labels, ds.labels_, atol=0.005)
     assert_series_equal(
         pd.Series([1 / 3, 2 / 3], pd.Index(["no", "yes"], name="label"), name="prior"),
         ds.priors_,