shape.py 449 B

123456789101112131415161718192021
  1. from typing import Any, NamedTuple, cast
  2. import numpy as np
  3. # Subtype of tuple[int, int]
  4. class XYGrid(NamedTuple):
  5. x_axis: int
  6. y_axis: int
  7. # TODO: remove this cast after: https://github.com/numpy/numpy/pull/27171
  8. arr: np.ndarray[XYGrid, Any] = cast(
  9. np.ndarray[XYGrid, Any],
  10. np.empty(XYGrid(2, 2)),
  11. )
  12. # Test variance of _ShapeType_co
  13. def accepts_2d(a: np.ndarray[tuple[int, int], Any]) -> None:
  14. return None
  15. accepts_2d(arr)