TYP: Add a protocol class representing `__array_function__`

2 files changed, 15 insertions, 0 deletions

diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
index d5cfbf5ac..72ac750ae 100644
--- a/numpy/typing/__init__.py
+++ b/numpy/typing/__init__.py
@@ -343,6 +343,7 @@ from ._array_like import (
     _ArrayLike,
     _FiniteNestedSequence,
     _SupportsArray,
+    _SupportsArrayFunc,
     _ArrayLikeInt,
     _ArrayLikeBool_co,
     _ArrayLikeUInt_co,
diff --git a/numpy/typing/_array_like.py b/numpy/typing/_array_like.py
index 02e5ee573..bba545b7b 100644
--- a/numpy/typing/_array_like.py
+++ b/numpy/typing/_array_like.py
@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+# NOTE: Import `Sequence` from `typing` as we it is needed for a type-alias,
+# not an annotation
+from collections.abc import Collection, Callable
 from typing import Any, Sequence, Protocol, Union, TypeVar
 from numpy import (
     ndarray,
@@ -34,6 +37,17 @@ class _SupportsArray(Protocol[_DType_co]):
     def __array__(self) -> ndarray[Any, _DType_co]: ...
 
 
+class _SupportsArrayFunc(Protocol):
+    """A protocol class representing `~class.__array_function__`."""
+    def __array_function__(
+        self,
+        func: Callable[..., Any],
+        types: Collection[type[Any]],
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> object: ...
+
+
 # TODO: Wait until mypy supports recursive objects in combination with typevars
 _FiniteNestedSequence = Union[
     _T,