More typing

Author: dustalov

python/evalica/__init__.py

@@ -6,7 +6,7 @@
 from collections.abc import Collection, Hashable
 from dataclasses import dataclass
 from types import MappingProxyType
-from typing import Generic, Literal, Protocol, TypeVar, runtime_checkable
+from typing import Generic, Literal, Protocol, TypeVar, cast, runtime_checkable
 
 import numpy as np
 import numpy.typing as npt
@@ -979,7 +979,7 @@ def pairwise_scores(
         raise ScoreDimensionError(scores.ndim)
 
     if solver == "naive":
-        return pairwise_scores_naive(scores)
+        return cast("npt.NDArray[np.float64]", pairwise_scores_naive(scores))
 
     return pairwise_scores_pyo3(scores)
 

python/evalica/naive.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, TypeVar, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -9,17 +9,19 @@
 
 if TYPE_CHECKING:
     from collections.abc import Collection
-    from typing import Any
 
+    S = TypeVar("S", bound=npt.NBitBase)
+    T = TypeVar("T")
 
-def pairwise_scores(scores: npt.NDArray[np.number[Any]]) -> npt.NDArray[np.float64]:
+
+def pairwise_scores(scores: npt.NDArray[np.number[S]]) -> npt.NDArray[np.number[S]]:
     if not scores.size:
-        return np.zeros((0, 0))
+        return np.zeros((0, 0), dtype=scores.dtype)
 
     return np.nan_to_num(scores[:, np.newaxis] / (scores + scores[:, np.newaxis]))
 
 
-def _check_lengths(xs: Collection[Any], *rest: Collection[Any]) -> None:
+def _check_lengths(xs: Collection[T], *rest: Collection[T]) -> None:
     length = len(xs)
 
     for collection in rest:
@@ -225,9 +227,9 @@ def eigen(
 
 
 def pagerank_matrix(
-        matrix: npt.NDArray[np.floating[Any]],
+        matrix: npt.NDArray[np.floating[S]],
         damping: float,
-) -> npt.NDArray[np.floating[Any]]:
+) -> npt.NDArray[np.floating[S]]:
     if not matrix.size:
         return np.zeros(0, dtype=matrix.dtype)