codekingpro/portable-devtools

1from collections.abc import Iterable
2from typing import Any, Literal, TypeAlias, TypeVar, overload
3
4from numpy._typing import DTypeLike, NDArray, _SupportsArrayFunc
5
6__all__ = ["require"]
7
8_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
9
10_Requirements: TypeAlias = Literal[
11    "C", "C_CONTIGUOUS", "CONTIGUOUS",
12    "F", "F_CONTIGUOUS", "FORTRAN",
13    "A", "ALIGNED",
14    "W", "WRITEABLE",
15    "O", "OWNDATA"
16]
17_E: TypeAlias = Literal["E", "ENSUREARRAY"]
18_RequirementsWithE: TypeAlias = _Requirements | _E
19
20@overload
21def require(
22    a: _ArrayT,
23    dtype: None = None,
24    requirements: _Requirements | Iterable[_Requirements] | None = None,
25    *,
26    like: _SupportsArrayFunc | None = None
27) -> _ArrayT: ...
28@overload
29def require(
30    a: object,
31    dtype: DTypeLike | None = None,
32    requirements: _E | Iterable[_RequirementsWithE] | None = None,
33    *,
34    like: _SupportsArrayFunc | None = None
35) -> NDArray[Any]: ...
36@overload
37def require(
38    a: object,
39    dtype: DTypeLike | None = None,
40    requirements: _Requirements | Iterable[_Requirements] | None = None,
41    *,
42    like: _SupportsArrayFunc | None = None
43) -> NDArray[Any]: ...
44