✨ stats: improved Covariance generic type parameter inference (#739)

Author: jorenham

scipy-stubs/stats/_covariance.pyi

@@ -1,5 +1,4 @@
-from collections.abc import Sequence
-from typing import Final, Generic, Protocol, TypeAlias, overload, type_check_only
+from typing import Final, Generic, Protocol, overload, type_check_only
 from typing_extensions import TypeVar
 
 import numpy as np
@@ -8,46 +7,49 @@ import optype.numpy.compat as npc
 
 __all__ = ["Covariance"]
 
-# `float16` and `longdouble` aren't supported in `scipy.linalg`, and neither is `bool_`
-_Scalar_uif: TypeAlias = np.float32 | np.float64 | npc.integer
+_ScalarT = TypeVar("_ScalarT", bound=npc.floating | npc.integer)
+_ScalarT_co = TypeVar("_ScalarT_co", bound=npc.floating | npc.integer, default=np.float64, covariant=True)
 
-_SCT = TypeVar("_SCT", bound=_Scalar_uif)
-_SCT_co = TypeVar("_SCT_co", bound=_Scalar_uif, covariant=True, default=np.float64)
-
-class Covariance(Generic[_SCT_co]):
+class Covariance(Generic[_ScalarT_co]):
     @staticmethod
     @overload
-    def from_diagonal(diagonal: Sequence[int]) -> CovViaDiagonal[np.int_]: ...
+    def from_diagonal(diagonal: onp.ToJustFloat64_1D) -> CovViaDiagonal[np.float64]: ...
     @staticmethod
     @overload
-    def from_diagonal(diagonal: Sequence[float]) -> CovViaDiagonal[np.int_ | np.float64]: ...
+    def from_diagonal(diagonal: onp.ToJustInt64_1D) -> CovViaDiagonal[np.int_]: ...
     @staticmethod
     @overload
-    def from_diagonal(diagonal: Sequence[_SCT] | onp.CanArrayND[_SCT]) -> CovViaDiagonal[_SCT]: ...
+    def from_diagonal(diagonal: onp.ToArray1D[_ScalarT, _ScalarT]) -> CovViaDiagonal[_ScalarT]: ...
+
+    #
     @staticmethod
     def from_precision(precision: onp.ToFloat2D, covariance: onp.ToFloat2D | None = None) -> CovViaPrecision: ...
     @staticmethod
     def from_cholesky(cholesky: onp.ToFloat2D) -> CovViaCholesky: ...
     @staticmethod
     def from_eigendecomposition(eigendecomposition: tuple[onp.ToFloat1D, onp.ToFloat2D]) -> CovViaEigendecomposition: ...
-    def whiten(self, /, x: onp.AnyIntegerArray | onp.AnyFloatingArray) -> onp.ArrayND[npc.floating]: ...
-    def colorize(self, /, x: onp.AnyIntegerArray | onp.AnyFloatingArray) -> onp.ArrayND[npc.floating]: ...
+
+    #
     @property
     def log_pdet(self, /) -> np.float64: ...
     @property
     def rank(self, /) -> np.int_: ...
     @property
-    def covariance(self, /) -> onp.Array2D[_SCT_co]: ...
+    def covariance(self, /) -> onp.Array2D[_ScalarT_co]: ...
     @property
     def shape(self, /) -> tuple[int, int]: ...
 
-class CovViaDiagonal(Covariance[_SCT_co], Generic[_SCT_co]):
+    #
+    def whiten(self, /, x: onp.ToFloatND) -> onp.ArrayND[npc.floating]: ...
+    def colorize(self, /, x: onp.ToFloatND) -> onp.ArrayND[npc.floating]: ...
+
+class CovViaDiagonal(Covariance[_ScalarT_co], Generic[_ScalarT_co]):
     @overload
-    def __init__(self: CovViaDiagonal[np.int_], /, diagonal: Sequence[int]) -> None: ...
+    def __init__(self: CovViaDiagonal[np.float64], /, diagonal: onp.ToJustFloat64_1D) -> None: ...
     @overload
-    def __init__(self: CovViaDiagonal[np.int_ | np.float64], /, diagonal: Sequence[float]) -> None: ...
+    def __init__(self: CovViaDiagonal[np.int_], /, diagonal: onp.ToJustInt64_1D) -> None: ...
     @overload
-    def __init__(self, /, diagonal: Sequence[float | _SCT_co] | onp.CanArrayND[_SCT_co]) -> None: ...
+    def __init__(self, /, diagonal: onp.ToArray1D[_ScalarT_co, _ScalarT_co]) -> None: ...
 
 class CovViaPrecision(Covariance[np.float64]):
     def __init__(self, /, precision: onp.ToFloat2D, covariance: onp.ToFloat2D | None = None) -> None: ...
@@ -63,17 +65,17 @@ class _PSD(Protocol):
     _M: onp.ArrayND[np.float64]
     V: onp.ArrayND[np.float64]
     U: onp.ArrayND[np.float64]
-    eps: np.float64 | float
-    log_pdet: np.float64 | float
-    cond: np.float64 | float
+    eps: float
+    log_pdet: float
+    cond: float
     rank: int
 
     @property
     def pinv(self, /) -> onp.ArrayND[npc.floating]: ...
 
 class CovViaPSD(Covariance[np.float64]):
     _LP: Final[onp.ArrayND[np.float64]]
-    _log_pdet: Final[np.float64 | float]
+    _log_pdet: Final[float]
     _rank: Final[int]
     _covariance: Final[onp.ArrayND[np.float64]]
     _shape: tuple[int, int]