🩹 add overload for ndarray.__matmul__

Author: guan404ming

src/numpy-stubs/__init__.pyi

@@ -27,6 +27,8 @@ from typing import (
 )
 from typing_extensions import Buffer, CapsuleType, LiteralString, Never, Protocol, Self, TypeVar, Unpack, deprecated, override
 
+import numpy as np
+
 from . import (
     __config__ as __config__,
     _array_api_info as _array_api_info,
@@ -611,6 +613,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)
 
@@ -2531,8 +2535,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     @overload
     def __imul__(self: NDArray[object_], rhs: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
 
-    # TODO(jorenham): Support the "1d @ 1d -> scalar" case
-    # https://github.com/numpy/numtype/issues/197
+    @overload
+    def __matmul__(self: _Array1D[_ScalarT], rhs: _Array1D[_ScalarT], /) -> _ScalarT: ...
     @overload
     def __matmul__(self: NDArray[_NumberT], rhs: _ArrayLikeBool_co, /) -> NDArray[_NumberT]: ...
     @overload
@@ -2566,12 +2570,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 +2610,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_]: ...