🩹 add overload for ndarray.__matmul__

Author: guan404ming

src/numpy-stubs/__init__.pyi

@@ -1,3 +1,4 @@
+import numpy as np
 import abc
 import ctypes as ct
 import datetime as dt
@@ -611,6 +612,8 @@ _DT64ItemT = TypeVar("_DT64ItemT", bound=dt.date | int | None)
 _DT64ItemT_co = TypeVar("_DT64ItemT_co", bound=dt.date | int | None, default=dt.date | int | None, covariant=True)
 _TD64UnitT = TypeVar("_TD64UnitT", bound=_TD64Unit, default=_TD64Unit)
 
+_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]]
+
 ###
 # Type Aliases (for internal use only)
 
@@ -2534,7 +2537,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     # TODO(jorenham): Support the "1d @ 1d -> scalar" case
     # https://github.com/numpy/numtype/issues/197
     @overload
-    def __matmul__(self: NDArray[_NumberT], rhs: _ArrayLikeBool_co, /) -> NDArray[_NumberT]: ...
+    def __matmul__(self: _Array1D[_ScalarT], rhs: _Array1D[_ScalarT], /) -> _ScalarT: ...
     @overload
     def __matmul__(self: NDArray[bool_], rhs: _ArrayLike[_NumberT], /) -> NDArray[_NumberT]: ...
     @overload
@@ -2566,12 +2569,14 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     @overload
     def __matmul__(self: NDArray[bool_ | number], rhs: _ArrayLikeNumber_co, /) -> NDArray[Incomplete]: ...
     @overload
-    def __matmul__(self: NDArray[object_], rhs: object, /) -> NDArray[object_]: ...
+    def __matmul__(self: NDArray[object_], rhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
     @overload
     def __matmul__(self, rhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
 
     # keep in sync with __matmul__
     @overload
+    def __rmatmul__(self: _Array1D[_ScalarT], rhs: _Array1D[_ScalarT], /) -> _ScalarT: ...
+    @overload
     def __rmatmul__(self: NDArray[_NumberT], lhs: _ArrayLikeBool_co, /) -> NDArray[_NumberT]: ...
     @overload
     def __rmatmul__(self: NDArray[bool_], lhs: _ArrayLike[_NumberT], /) -> NDArray[_NumberT]: ...
@@ -2604,7 +2609,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     @overload
     def __rmatmul__(self: NDArray[bool_ | number], lhs: _ArrayLikeNumber_co, /) -> NDArray[Incomplete]: ...
     @overload
-    def __rmatmul__(self: NDArray[object_], lhs: object, /) -> NDArray[object_]: ...
+    def __rmatmul__(self: NDArray[object_], lhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
     @overload
     def __rmatmul__(self, lhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...