TYP: Add a protocol class representing __array_function__

BvB93 · BvB93 · commit d2df8c6a831a · 2022-01-11T18:48:44.000+01:00
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
@@ -343,6 +343,7 @@ class _8Bit(_16Bit):  # type: ignore[misc]
     _ArrayLike,
     _FiniteNestedSequence,
     _SupportsArray,
+    _SupportsArrayFunc,
     _ArrayLikeInt,
     _ArrayLikeBool_co,
     _ArrayLikeUInt_co,
diff --git 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,
