summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
authorBas van Beek <b.f.van.beek@vu.nl>2021-09-18 12:49:57 +0200
committerBas van Beek <b.f.van.beek@vu.nl>2021-09-18 12:49:57 +0200
commitf8958a2305c88f1c97b2cd5256d7e29faca93f13 (patch)
tree4c73aa61847f9f665daffeee0f7b50c5ee498cf9 /numpy
parentd5cee25a90fa651c9e865c72de872b9a4a0a8016 (diff)
downloadnumpy-f8958a2305c88f1c97b2cd5256d7e29faca93f13.tar.gz
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 <bashtage@users.noreply.github.com>
Diffstat (limited to 'numpy')
-rw-r--r--numpy/linalg/linalg.pyi18
-rw-r--r--numpy/typing/tests/data/reveal/linalg.py8
2 files changed, 17 insertions, 9 deletions
diff --git a/numpy/linalg/linalg.pyi b/numpy/linalg/linalg.pyi
index cd9fd1493..3ce0e2229 100644
--- a/numpy/linalg/linalg.pyi
+++ b/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]]]: ...
diff --git a/numpy/typing/tests/data/reveal/linalg.py b/numpy/typing/tests/data/reveal/linalg.py
index f4f52d085..fecdc0d37 100644
--- a/numpy/typing/tests/data/reveal/linalg.py
+++ b/numpy/typing/tests/data/reveal/linalg.py
@@ -37,16 +37,16 @@ reveal_type(np.linalg.qr(AR_i8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{flo
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}]]]