from collections.abc import Sequence

from typing import Any, NotRequired, Protocol, TypedDict, runtime_checkable

import numpy as np

from ._char_codes import (

_BoolCodes,

_BytesCodes,

_ComplexFloatingCodes,

_DT64Codes,

_FloatingCodes,

_NumberCodes,

_ObjectCodes,

_SignedIntegerCodes,

_StrCodes,

_TD64Codes,

_UnsignedIntegerCodes,

_VoidCodes,

)

type _DTypeLikeNested = Any # TODO: wait for support for recursive types

class _DTypeDict(TypedDict):

names: Sequence[str]

formats: Sequence[_DTypeLikeNested]

# Only `str` elements are usable as indexing aliases,

# but `titles` can in principle accept any object

offsets: NotRequired[Sequence[int]]

titles: NotRequired[Sequence[Any]]

itemsize: NotRequired[int]

aligned: NotRequired[bool]

# A protocol for anything with the dtype attribute

@runtime_checkable

class _HasDType[DTypeT: np.dtype](Protocol):

@property

def dtype(self) -> DTypeT: ...

@runtime_checkable

class _HasNumPyDType[DTypeT: np.dtype](Protocol):

@property

def __numpy_dtype__(self, /) -> DTypeT: ...

type _SupportsDType[DTypeT: np.dtype] = _HasDType[DTypeT] | _HasNumPyDType[DTypeT]

# A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic`

type _DTypeLike[ScalarT: np.generic] = (

type[ScalarT] | np.dtype[ScalarT] | _SupportsDType[np.dtype[ScalarT]]

)

# Would create a dtype[np.void]

type _VoidDTypeLike = (

# If a tuple, then it can be either:

# - (flexible_dtype, itemsize)

# - (fixed_dtype, shape)

# - (base_dtype, new_dtype)

# But because `_DTypeLikeNested = Any`, the first two cases are redundant

# tuple[_DTypeLikeNested, int] | tuple[_DTypeLikeNested, _ShapeLike] |

tuple[_DTypeLikeNested, _DTypeLikeNested]

# [(field_name, field_dtype, field_shape), ...]

# The type here is quite broad because NumPy accepts quite a wide

# range of inputs inside the list; see the tests for some examples.

| list[Any]

# {'names': ..., 'formats': ..., 'offsets': ..., 'titles': ..., 'itemsize': ...}

| _DTypeDict

)

# Aliases for commonly used dtype-like objects.

# Note that the precision of `np.number` subclasses is ignored herein.

type _DTypeLikeBool = type[bool] | _DTypeLike[np.bool] | _BoolCodes

type _DTypeLikeInt = type[int] | _DTypeLike[np.signedinteger] | _SignedIntegerCodes

type _DTypeLikeUInt = _DTypeLike[np.unsignedinteger] | _UnsignedIntegerCodes

type _DTypeLikeFloat = type[float] | _DTypeLike[np.floating] | _FloatingCodes

type _DTypeLikeComplex = (

type[complex] | _DTypeLike[np.complexfloating] | _ComplexFloatingCodes

)

type _DTypeLikeComplex_co = (

type[complex] | _DTypeLike[np.bool | np.number] | _BoolCodes | _NumberCodes

)

type _DTypeLikeDT64 = _DTypeLike[np.timedelta64] | _TD64Codes

type _DTypeLikeTD64 = _DTypeLike[np.datetime64] | _DT64Codes

type _DTypeLikeBytes = type[bytes] | _DTypeLike[np.bytes_] | _BytesCodes

type _DTypeLikeStr = type[str] | _DTypeLike[np.str_] | _StrCodes

type _DTypeLikeVoid = (

type[memoryview] | _DTypeLike[np.void] | _VoidDTypeLike | _VoidCodes

)

type _DTypeLikeObject = type[object] | _DTypeLike[np.object_] | _ObjectCodes

# Anything that can be coerced into numpy.dtype.

# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html

type DTypeLike = type | str | np.dtype | _SupportsDType[np.dtype] | _VoidDTypeLike

# NOTE: while it is possible to provide the dtype as a dict of

# dtype-like objects (e.g. `{'field1': ..., 'field2': ..., ...}`),

# this syntax is officially discouraged and

# therefore not included in the type-union defining `DTypeLike`.

#

# See https://github.com/numpy/numpy/issues/16891 for more details.