_typing.py 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. """
  2. This file defines the types for type annotations.
  3. These names aren't part of the module namespace, but they are used in the
  4. annotations in the function signatures. The functions in the module are only
  5. valid for inputs that match the given type annotations.
  6. """
  7. from __future__ import annotations
  8. __all__ = [
  9. "Array",
  10. "Device",
  11. "Dtype",
  12. "SupportsDLPack",
  13. "SupportsBufferProtocol",
  14. "PyCapsule",
  15. ]
  16. import sys
  17. from typing import (
  18. Any,
  19. Literal,
  20. Sequence,
  21. Type,
  22. Union,
  23. TYPE_CHECKING,
  24. TypeVar,
  25. Protocol,
  26. )
  27. from ._array_object import Array
  28. from numpy import (
  29. dtype,
  30. int8,
  31. int16,
  32. int32,
  33. int64,
  34. uint8,
  35. uint16,
  36. uint32,
  37. uint64,
  38. float32,
  39. float64,
  40. )
  41. _T_co = TypeVar("_T_co", covariant=True)
  42. class NestedSequence(Protocol[_T_co]):
  43. def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
  44. def __len__(self, /) -> int: ...
  45. Device = Literal["cpu"]
  46. if TYPE_CHECKING or sys.version_info >= (3, 9):
  47. Dtype = dtype[Union[
  48. int8,
  49. int16,
  50. int32,
  51. int64,
  52. uint8,
  53. uint16,
  54. uint32,
  55. uint64,
  56. float32,
  57. float64,
  58. ]]
  59. else:
  60. Dtype = dtype
  61. SupportsBufferProtocol = Any
  62. PyCapsule = Any
  63. class SupportsDLPack(Protocol):
  64. def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...