MAINT: Let eig and eigvals return a union for real inputs

Whether an array of a `floating` or `complexfloating` dtype is returned depends on the values within in the supplied matrix.

Co-Authored-By: Kevin Sheppard <[email protected]>

Author: Bas van Beek

numpy/linalg/linalg.pyi

@@ -10,7 +10,15 @@ from typing import (
     Tuple,
 )
 
-from numpy import generic, int32, floating, float64, complexfloating
+from numpy import (
+    generic,
+    floating,
+    complexfloating,
+    int32,
+    float64,
+    complex128,
+)
+
 from numpy.typing import (
     NDArray,
     ArrayLike,
@@ -109,9 +117,9 @@ def qr(a: _ArrayLikeFloat_co, mode: _ModeKind = ...) -> _2Tuple[NDArray[floating
 def qr(a: _ArrayLikeComplex_co, mode: _ModeKind = ...) -> _2Tuple[NDArray[complexfloating[Any, Any]]]: ...
 
 @overload
-def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64]: ...
+def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64] | NDArray[complex128]: ...
 @overload
-def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ...
+def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]] | NDArray[complexfloating[Any, Any]]: ...
 @overload
 def eigvals(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
 
@@ -121,9 +129,9 @@ def eigvalsh(a: _ArrayLikeInt_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[
 def eigvalsh(a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[floating[Any]]: ...
 
 @overload
-def eig(a: _ArrayLikeInt_co) -> _2Tuple[NDArray[float64]]: ...
+def eig(a: _ArrayLikeInt_co) -> _2Tuple[NDArray[float64]] | _2Tuple[NDArray[complex128]]: ...
 @overload
-def eig(a: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ...
+def eig(a: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]] | _2Tuple[NDArray[complexfloating[Any, Any]]]: ...
 @overload
 def eig(a: _ArrayLikeComplex_co) -> _2Tuple[NDArray[complexfloating[Any, Any]]]: ...
 

numpy/typing/tests/data/reveal/linalg.py

@@ -37,16 +37,16 @@
 reveal_type(np.linalg.qr(AR_f8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
 reveal_type(np.linalg.qr(AR_c16))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]]
 
-reveal_type(np.linalg.eigvals(AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
-reveal_type(np.linalg.eigvals(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.linalg.eigvals(AR_i8))  # E: Union[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{complex128}]]]
+reveal_type(np.linalg.eigvals(AR_f8))  # E: Union[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]]
 reveal_type(np.linalg.eigvals(AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
 
 reveal_type(np.linalg.eigvalsh(AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
 reveal_type(np.linalg.eigvalsh(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
 reveal_type(np.linalg.eigvalsh(AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
 
-reveal_type(np.linalg.eig(AR_i8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]]
-reveal_type(np.linalg.eig(AR_f8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
+reveal_type(np.linalg.eig(AR_i8))  # E: Union[Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]], Tuple[numpy.ndarray[Any, numpy.dtype[{complex128}]], numpy.ndarray[Any, numpy.dtype[{complex128}]]]]
+reveal_type(np.linalg.eig(AR_f8))  # E: Union[Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]], Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]]]
 reveal_type(np.linalg.eig(AR_c16))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]]
 
 reveal_type(np.linalg.eigh(AR_i8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]]