Skip to content

Commit 6c38e05

Browse files
charrischarris
authored
Merge pull request numpy#20105 from BvB93/matrix
ENH: Add annotations for `np.matrix`
2 parents 8a15d8a + 228cd06 commit 6c38e05

File tree

4 files changed

+211
-52
lines changed

4 files changed

+211
-52
lines changed

‎numpy/__init__.pyi‎

Lines changed: 120 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -724,53 +724,6 @@ class chararray(ndarray[_ShapeType, _DType_co]):
724724
def isnumeric(self): ...
725725
def isdecimal(self): ...
726726

727-
class matrix(ndarray[_ShapeType, _DType_co]):
728-
def __new__(
729-
subtype,
730-
data: Any,
731-
dtype: Any = ...,
732-
copy: Any = ...,
733-
) -> Any: ...
734-
def __array_finalize__(self, obj): ...
735-
def __getitem__(self, index): ...
736-
def __mul__(self, other): ...
737-
def __rmul__(self, other): ...
738-
def __imul__(self, other): ...
739-
def __pow__(self, other): ...
740-
def __ipow__(self, other): ...
741-
def __rpow__(self, other): ...
742-
def tolist(self): ...
743-
def sum(self, axis=..., dtype=..., out=...): ...
744-
def squeeze(self, axis=...): ...
745-
def flatten(self, order=...): ...
746-
def mean(self, axis=..., dtype=..., out=...): ...
747-
def std(self, axis=..., dtype=..., out=..., ddof=...): ...
748-
def var(self, axis=..., dtype=..., out=..., ddof=...): ...
749-
def prod(self, axis=..., dtype=..., out=...): ...
750-
def any(self, axis=..., out=...): ...
751-
def all(self, axis=..., out=...): ...
752-
def max(self, axis=..., out=...): ...
753-
def argmax(self, axis=..., out=...): ...
754-
def min(self, axis=..., out=...): ...
755-
def argmin(self, axis=..., out=...): ...
756-
def ptp(self, axis=..., out=...): ...
757-
def ravel(self, order=...): ...
758-
@property
759-
def T(self): ...
760-
@property
761-
def I(self): ...
762-
@property
763-
def A(self): ...
764-
@property
765-
def A1(self): ...
766-
@property
767-
def H(self): ...
768-
def getT(self): ...
769-
def getA(self): ...
770-
def getA1(self): ...
771-
def getH(self): ...
772-
def getI(self): ...
773-
774727
# Some of these are aliases; others are wrappers with an identical signature
775728
round = around
776729
round_ = around
@@ -3916,3 +3869,123 @@ class poly1d:
39163869
m: SupportsInt | SupportsIndex = ...,
39173870
k: None | _ArrayLikeComplex_co | _ArrayLikeObject_co = ...,
39183871
) -> poly1d: ...
3872+
3873+
class matrix(ndarray[_ShapeType, _DType_co]):
3874+
def __new__(
3875+
subtype,
3876+
data: ArrayLike,
3877+
dtype: DTypeLike = ...,
3878+
copy: bool = ...,
3879+
) -> matrix[Any, Any]: ...
3880+
def __array_finalize__(self, obj: NDArray[Any]) -> None: ...
3881+
def __getitem__(self, index): ... # TODO
3882+
def __mul__(self, other: ArrayLike) -> matrix[Any, Any]: ...
3883+
def __rmul__(self, other: ArrayLike) -> matrix[Any, Any]: ...
3884+
def __imul__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ...
3885+
def __pow__(self, other: ArrayLike) -> matrix[Any, Any]: ...
3886+
def __ipow__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ...
3887+
3888+
@overload
3889+
def sum(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ...
3890+
@overload
3891+
def sum(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ...
3892+
@overload
3893+
def sum(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3894+
3895+
@overload
3896+
def mean(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ...
3897+
@overload
3898+
def mean(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ...
3899+
@overload
3900+
def mean(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3901+
3902+
@overload
3903+
def std(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ...
3904+
@overload
3905+
def std(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ...
3906+
@overload
3907+
def std(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ...
3908+
3909+
@overload
3910+
def var(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ...
3911+
@overload
3912+
def var(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ...
3913+
@overload
3914+
def var(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ...
3915+
3916+
@overload
3917+
def prod(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ...
3918+
@overload
3919+
def prod(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ...
3920+
@overload
3921+
def prod(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3922+
3923+
@overload
3924+
def any(self, axis: None = ..., out: None = ...) -> bool_: ...
3925+
@overload
3926+
def any(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ...
3927+
@overload
3928+
def any(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3929+
3930+
@overload
3931+
def all(self, axis: None = ..., out: None = ...) -> bool_: ...
3932+
@overload
3933+
def all(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ...
3934+
@overload
3935+
def all(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3936+
3937+
@overload
3938+
def max(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ...
3939+
@overload
3940+
def max(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ...
3941+
@overload
3942+
def max(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3943+
3944+
@overload
3945+
def min(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ...
3946+
@overload
3947+
def min(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ...
3948+
@overload
3949+
def min(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3950+
3951+
@overload
3952+
def argmax(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ...
3953+
@overload
3954+
def argmax(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ...
3955+
@overload
3956+
def argmax(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3957+
3958+
@overload
3959+
def argmin(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ...
3960+
@overload
3961+
def argmin(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ...
3962+
@overload
3963+
def argmin(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3964+
3965+
@overload
3966+
def ptp(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ...
3967+
@overload
3968+
def ptp(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ...
3969+
@overload
3970+
def ptp(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ...
3971+
3972+
def squeeze(self, axis: None | _ShapeLike = ...) -> matrix[Any, _DType_co]: ...
3973+
def tolist(self: matrix[Any, dtype[_SupportsItem[_T]]]) -> List[List[_T]]: ... # type: ignore[typevar]
3974+
def ravel(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ...
3975+
def flatten(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ...
3976+
3977+
@property
3978+
def T(self) -> matrix[Any, _DType_co]: ...
3979+
@property
3980+
def I(self) -> matrix[Any, Any]: ...
3981+
@property
3982+
def A(self) -> ndarray[_ShapeType, _DType_co]: ...
3983+
@property
3984+
def A1(self) -> ndarray[Any, _DType_co]: ...
3985+
@property
3986+
def H(self) -> matrix[Any, _DType_co]: ...
3987+
def getT(self) -> matrix[Any, _DType_co]: ...
3988+
def getI(self) -> matrix[Any, Any]: ...
3989+
def getA(self) -> ndarray[_ShapeType, _DType_co]: ...
3990+
def getA1(self) -> ndarray[Any, _DType_co]: ...
3991+
def getH(self) -> matrix[Any, _DType_co]: ...

‎numpy/matrixlib/__init__.pyi‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
1-
from typing import Any, List
1+
from typing import List
22

33
from numpy._pytesttester import PytestTester
44

55
from numpy import (
66
matrix as matrix,
77
)
88

9+
from numpy.matrixlib.defmatrix import (
10+
bmat as bmat,
11+
mat as mat,
12+
asmatrix as asmatrix,
13+
)
14+
915
__all__: List[str]
1016
__path__: List[str]
1117
test: PytestTester
12-
13-
def bmat(obj, ldict=..., gdict=...): ...
14-
def asmatrix(data, dtype=...): ...
15-
mat = asmatrix

‎numpy/matrixlib/defmatrix.pyi‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
from typing import List, Any, Sequence, Mapping
2+
from numpy import matrix as matrix
3+
from numpy.typing import ArrayLike, DTypeLike, NDArray
4+
5+
__all__: List[str]
6+
7+
def bmat(
8+
obj: str | Sequence[ArrayLike] | NDArray[Any],
9+
ldict: None | Mapping[str, Any] = ...,
10+
gdict: None | Mapping[str, Any] = ...,
11+
) -> matrix[Any, Any]: ...
12+
13+
def asmatrix(data: ArrayLike, dtype: DTypeLike = ...) -> matrix[Any, Any]: ...
14+
15+
mat = asmatrix

‎numpy/typing/tests/data/reveal/matrix.pyi‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
from typing import Any
2+
import numpy as np
3+
import numpy.typing as npt
4+
5+
mat: np.matrix[Any, np.dtype[np.int64]]
6+
ar_f8: npt.NDArray[np.float64]
7+
8+
reveal_type(mat * 5) # E: numpy.matrix[Any, Any]
9+
reveal_type(5 * mat) # E: numpy.matrix[Any, Any]
10+
mat *= 5
11+
12+
reveal_type(mat**5) # E: numpy.matrix[Any, Any]
13+
mat **= 5
14+
15+
reveal_type(mat.sum()) # E: Any
16+
reveal_type(mat.mean()) # E: Any
17+
reveal_type(mat.std()) # E: Any
18+
reveal_type(mat.var()) # E: Any
19+
reveal_type(mat.prod()) # E: Any
20+
reveal_type(mat.any()) # E: numpy.bool_
21+
reveal_type(mat.all()) # E: numpy.bool_
22+
reveal_type(mat.max()) # E: {int64}
23+
reveal_type(mat.min()) # E: {int64}
24+
reveal_type(mat.argmax()) # E: {intp}
25+
reveal_type(mat.argmin()) # E: {intp}
26+
reveal_type(mat.ptp()) # E: {int64}
27+
28+
reveal_type(mat.sum(axis=0)) # E: numpy.matrix[Any, Any]
29+
reveal_type(mat.mean(axis=0)) # E: numpy.matrix[Any, Any]
30+
reveal_type(mat.std(axis=0)) # E: numpy.matrix[Any, Any]
31+
reveal_type(mat.var(axis=0)) # E: numpy.matrix[Any, Any]
32+
reveal_type(mat.prod(axis=0)) # E: numpy.matrix[Any, Any]
33+
reveal_type(mat.any(axis=0)) # E: numpy.matrix[Any, numpy.dtype[numpy.bool_]]
34+
reveal_type(mat.all(axis=0)) # E: numpy.matrix[Any, numpy.dtype[numpy.bool_]]
35+
reveal_type(mat.max(axis=0)) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
36+
reveal_type(mat.min(axis=0)) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
37+
reveal_type(mat.argmax(axis=0)) # E: numpy.matrix[Any, numpy.dtype[{intp}]]
38+
reveal_type(mat.argmin(axis=0)) # E: numpy.matrix[Any, numpy.dtype[{intp}]]
39+
reveal_type(mat.ptp(axis=0)) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
40+
41+
reveal_type(mat.sum(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
42+
reveal_type(mat.mean(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
43+
reveal_type(mat.std(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
44+
reveal_type(mat.var(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
45+
reveal_type(mat.prod(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
46+
reveal_type(mat.any(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
47+
reveal_type(mat.all(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
48+
reveal_type(mat.max(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
49+
reveal_type(mat.min(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
50+
reveal_type(mat.argmax(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
51+
reveal_type(mat.argmin(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
52+
reveal_type(mat.ptp(out=ar_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
53+
54+
reveal_type(mat.T) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
55+
reveal_type(mat.I) # E: numpy.matrix[Any, Any]
56+
reveal_type(mat.A) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
57+
reveal_type(mat.A1) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
58+
reveal_type(mat.H) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
59+
reveal_type(mat.getT()) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
60+
reveal_type(mat.getI()) # E: numpy.matrix[Any, Any]
61+
reveal_type(mat.getA()) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
62+
reveal_type(mat.getA1()) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
63+
reveal_type(mat.getH()) # E: numpy.matrix[Any, numpy.dtype[{int64}]]
64+
65+
reveal_type(np.bmat(ar_f8)) # E: numpy.matrix[Any, Any]
66+
reveal_type(np.bmat([[0, 1, 2]])) # E: numpy.matrix[Any, Any]
67+
reveal_type(np.bmat("mat")) # E: numpy.matrix[Any, Any]
68+
69+
reveal_type(np.asmatrix(ar_f8, dtype=np.int64)) # E: numpy.matrix[Any, Any]

0 commit comments

Comments
 (0)