_array_like.py 1021 B

12345678910111213141516171819202122232425262728293031323334353637383940
  1. import sys
  2. from typing import Any, overload, Sequence, TYPE_CHECKING, Union
  3. from numpy import ndarray
  4. from ._scalars import _ScalarLike
  5. from ._dtype_like import DTypeLike
  6. if sys.version_info >= (3, 8):
  7. from typing import Protocol
  8. HAVE_PROTOCOL = True
  9. else:
  10. try:
  11. from typing_extensions import Protocol
  12. except ImportError:
  13. HAVE_PROTOCOL = False
  14. else:
  15. HAVE_PROTOCOL = True
  16. if TYPE_CHECKING or HAVE_PROTOCOL:
  17. class _SupportsArray(Protocol):
  18. @overload
  19. def __array__(self, __dtype: DTypeLike = ...) -> ndarray: ...
  20. @overload
  21. def __array__(self, dtype: DTypeLike = ...) -> ndarray: ...
  22. else:
  23. _SupportsArray = Any
  24. # TODO: support buffer protocols once
  25. #
  26. # https://bugs.python.org/issue27501
  27. #
  28. # is resolved. See also the mypy issue:
  29. #
  30. # https://github.com/python/typing/issues/593
  31. ArrayLike = Union[
  32. _ScalarLike,
  33. Sequence[_ScalarLike],
  34. Sequence[Sequence[Any]], # TODO: Wait for support for recursive types
  35. _SupportsArray,
  36. ]