437 lines
13 KiB
Python
437 lines
13 KiB
Python
from _typeshed import Incomplete
|
|
from collections.abc import Callable, Iterable, Mapping, Sequence
|
|
from typing import Any, Literal, overload
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
from numpy import _CastingKind
|
|
from numpy._typing import _AnyShape, _DTypeLike, _DTypeLikeVoid, _Shape
|
|
from numpy.ma.mrecords import MaskedRecords
|
|
|
|
__all__ = [
|
|
"append_fields",
|
|
"apply_along_fields",
|
|
"assign_fields_by_name",
|
|
"drop_fields",
|
|
"find_duplicates",
|
|
"flatten_descr",
|
|
"get_fieldstructure",
|
|
"get_names",
|
|
"get_names_flat",
|
|
"join_by",
|
|
"merge_arrays",
|
|
"rec_append_fields",
|
|
"rec_drop_fields",
|
|
"rec_join",
|
|
"recursive_fill_fields",
|
|
"rename_fields",
|
|
"repack_fields",
|
|
"require_fields",
|
|
"stack_arrays",
|
|
"structured_to_unstructured",
|
|
"unstructured_to_structured",
|
|
]
|
|
|
|
type _OneOrMany[T] = T | Iterable[T]
|
|
type _BuiltinSequence[T] = tuple[T, ...] | list[T]
|
|
|
|
type _NestedNames = tuple[str | _NestedNames, ...]
|
|
type _NonVoid = np.bool | np.number | np.character | np.datetime64 | np.timedelta64 | np.object_
|
|
type _NonVoidDType = np.dtype[_NonVoid] | np.dtypes.StringDType
|
|
|
|
type _JoinType = Literal["inner", "outer", "leftouter"]
|
|
|
|
###
|
|
|
|
def recursive_fill_fields[VoidArrayT: npt.NDArray[np.void]](input: npt.NDArray[np.void], output: VoidArrayT) -> VoidArrayT: ...
|
|
|
|
#
|
|
def get_names(adtype: np.dtype[np.void]) -> _NestedNames: ...
|
|
def get_names_flat(adtype: np.dtype[np.void]) -> tuple[str, ...]: ...
|
|
|
|
#
|
|
@overload
|
|
def flatten_descr[NonVoidDTypeT: _NonVoidDType](ndtype: NonVoidDTypeT) -> tuple[tuple[Literal[""], NonVoidDTypeT]]: ...
|
|
@overload
|
|
def flatten_descr(ndtype: np.dtype[np.void]) -> tuple[tuple[str, np.dtype]]: ...
|
|
|
|
#
|
|
def get_fieldstructure(
|
|
adtype: np.dtype[np.void],
|
|
lastname: str | None = None,
|
|
parents: dict[str, list[str]] | None = None,
|
|
) -> dict[str, list[str]]: ...
|
|
|
|
#
|
|
@overload
|
|
def merge_arrays[ShapeT: _Shape](
|
|
seqarrays: Sequence[np.ndarray[ShapeT, np.dtype]] | np.ndarray[ShapeT, np.dtype],
|
|
fill_value: float = -1,
|
|
flatten: bool = False,
|
|
usemask: bool = False,
|
|
asrecarray: bool = False,
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def merge_arrays(
|
|
seqarrays: Sequence[npt.ArrayLike] | np.void,
|
|
fill_value: float = -1,
|
|
flatten: bool = False,
|
|
usemask: bool = False,
|
|
asrecarray: bool = False,
|
|
) -> np.recarray[_AnyShape, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
@overload
|
|
def drop_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
drop_names: str | Iterable[str],
|
|
usemask: bool = True,
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def drop_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
drop_names: str | Iterable[str],
|
|
usemask: bool,
|
|
asrecarray: Literal[True],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def drop_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
drop_names: str | Iterable[str],
|
|
usemask: bool = True,
|
|
*,
|
|
asrecarray: Literal[True],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
@overload
|
|
def rename_fields[ShapeT: _Shape](
|
|
base: MaskedRecords[ShapeT, np.dtype[np.void]],
|
|
namemapper: Mapping[str, str],
|
|
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def rename_fields[ShapeT: _Shape](
|
|
base: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
|
|
namemapper: Mapping[str, str],
|
|
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def rename_fields[ShapeT: _Shape](
|
|
base: np.recarray[ShapeT, np.dtype[np.void]],
|
|
namemapper: Mapping[str, str],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def rename_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
namemapper: Mapping[str, str],
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None,
|
|
fill_value: int,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None = None,
|
|
fill_value: int = -1,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None,
|
|
fill_value: int,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[True],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None = None,
|
|
fill_value: int = -1,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[True],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None = None,
|
|
fill_value: int = -1,
|
|
usemask: Literal[True] = True,
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None,
|
|
fill_value: int,
|
|
usemask: Literal[True],
|
|
asrecarray: Literal[True],
|
|
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None = None,
|
|
fill_value: int = -1,
|
|
usemask: Literal[True] = True,
|
|
*,
|
|
asrecarray: Literal[True],
|
|
) -> MaskedRecords[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
def rec_drop_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
drop_names: str | Iterable[str],
|
|
) -> np.recarray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
def rec_append_fields[ShapeT: _Shape](
|
|
base: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
names: _OneOrMany[str],
|
|
data: _OneOrMany[npt.NDArray[Any]],
|
|
dtypes: _BuiltinSequence[np.dtype] | None = None,
|
|
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
# TODO(jorenham): Stop passing `void` directly once structured dtypes are implemented,
|
|
# e.g. using a `TypeVar` with constraints.
|
|
# https://github.com/numpy/numtype/issues/92
|
|
@overload
|
|
def repack_fields[DTypeT: np.dtype](a: DTypeT, align: bool = False, recurse: bool = False) -> DTypeT: ...
|
|
@overload
|
|
def repack_fields[ScalarT: np.generic](a: ScalarT, align: bool = False, recurse: bool = False) -> ScalarT: ...
|
|
@overload
|
|
def repack_fields[ArrayT: np.ndarray](a: ArrayT, align: bool = False, recurse: bool = False) -> ArrayT: ...
|
|
|
|
# TODO(jorenham): Attempt shape-typing (return type has ndim == arr.ndim + 1)
|
|
@overload
|
|
def structured_to_unstructured[ScalarT: np.generic](
|
|
arr: npt.NDArray[np.void],
|
|
dtype: _DTypeLike[ScalarT],
|
|
copy: bool = False,
|
|
casting: _CastingKind = "unsafe",
|
|
) -> npt.NDArray[ScalarT]: ...
|
|
@overload
|
|
def structured_to_unstructured(
|
|
arr: npt.NDArray[np.void],
|
|
dtype: npt.DTypeLike | None = None,
|
|
copy: bool = False,
|
|
casting: _CastingKind = "unsafe",
|
|
) -> npt.NDArray[Any]: ...
|
|
|
|
#
|
|
@overload
|
|
def unstructured_to_structured(
|
|
arr: npt.NDArray[Any],
|
|
dtype: npt.DTypeLike,
|
|
names: None = None,
|
|
align: bool = False,
|
|
copy: bool = False,
|
|
casting: str = "unsafe",
|
|
) -> npt.NDArray[np.void]: ...
|
|
@overload
|
|
def unstructured_to_structured(
|
|
arr: npt.NDArray[Any],
|
|
dtype: None,
|
|
names: _OneOrMany[str],
|
|
align: bool = False,
|
|
copy: bool = False,
|
|
casting: str = "unsafe",
|
|
) -> npt.NDArray[np.void]: ...
|
|
@overload
|
|
def unstructured_to_structured(
|
|
arr: npt.NDArray[Any],
|
|
dtype: None = None,
|
|
*,
|
|
names: _OneOrMany[str],
|
|
align: bool = False,
|
|
copy: bool = False,
|
|
casting: str = "unsafe",
|
|
) -> npt.NDArray[np.void]: ...
|
|
|
|
#
|
|
def apply_along_fields[ShapeT: _Shape](
|
|
func: Callable[[np.ndarray[ShapeT]], np.ndarray],
|
|
arr: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
def assign_fields_by_name(dst: npt.NDArray[np.void], src: npt.NDArray[np.void], zero_unassigned: bool = True) -> None: ...
|
|
|
|
#
|
|
def require_fields[ShapeT: _Shape](
|
|
array: np.ndarray[ShapeT, np.dtype[np.void]],
|
|
required_dtype: _DTypeLikeVoid,
|
|
) -> np.ndarray[ShapeT, np.dtype[np.void]]: ...
|
|
|
|
# TODO(jorenham): Attempt shape-typing
|
|
@overload
|
|
def stack_arrays[ArrayT: np.ndarray](
|
|
arrays: ArrayT,
|
|
defaults: Mapping[str, object] | None = None,
|
|
usemask: bool = True,
|
|
asrecarray: bool = False,
|
|
autoconvert: bool = False,
|
|
) -> ArrayT: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[False] = False,
|
|
autoconvert: bool = False,
|
|
) -> npt.NDArray[np.void]: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None = None,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[False] = False,
|
|
autoconvert: bool = False,
|
|
) -> npt.NDArray[np.void]: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None = None,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[True],
|
|
autoconvert: bool = False,
|
|
) -> np.recarray[_AnyShape, np.dtype[np.void]]: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None = None,
|
|
usemask: Literal[True] = True,
|
|
asrecarray: Literal[False] = False,
|
|
autoconvert: bool = False,
|
|
) -> np.ma.MaskedArray[_AnyShape, np.dtype[np.void]]: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None,
|
|
usemask: Literal[True],
|
|
asrecarray: Literal[True],
|
|
autoconvert: bool = False,
|
|
) -> MaskedRecords[_AnyShape, np.dtype[np.void]]: ...
|
|
@overload
|
|
def stack_arrays(
|
|
arrays: Sequence[npt.NDArray[Any]],
|
|
defaults: Mapping[str, Incomplete] | None = None,
|
|
usemask: Literal[True] = True,
|
|
*,
|
|
asrecarray: Literal[True],
|
|
autoconvert: bool = False,
|
|
) -> MaskedRecords[_AnyShape, np.dtype[np.void]]: ...
|
|
|
|
#
|
|
@overload
|
|
def find_duplicates[ShapeT: _Shape](
|
|
a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
|
|
key: str | None = None,
|
|
ignoremask: bool = True,
|
|
return_index: Literal[False] = False,
|
|
) -> np.ma.MaskedArray[ShapeT, np.dtype[np.void]]: ...
|
|
@overload
|
|
def find_duplicates[ShapeT: _Shape](
|
|
a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
|
|
key: str | None,
|
|
ignoremask: bool,
|
|
return_index: Literal[True],
|
|
) -> tuple[np.ma.MaskedArray[ShapeT, np.dtype[np.void]], np.ndarray[ShapeT, np.dtype[np.int_]]]: ...
|
|
@overload
|
|
def find_duplicates[ShapeT: _Shape](
|
|
a: np.ma.MaskedArray[ShapeT, np.dtype[np.void]],
|
|
key: str | None = None,
|
|
ignoremask: bool = True,
|
|
*,
|
|
return_index: Literal[True],
|
|
) -> tuple[np.ma.MaskedArray[ShapeT, np.dtype[np.void]], np.ndarray[ShapeT, np.dtype[np.int_]]]: ...
|
|
|
|
#
|
|
@overload
|
|
def join_by(
|
|
key: str | Sequence[str],
|
|
r1: npt.NDArray[np.void],
|
|
r2: npt.NDArray[np.void],
|
|
jointype: _JoinType = "inner",
|
|
r1postfix: str = "1",
|
|
r2postfix: str = "2",
|
|
defaults: Mapping[str, object] | None = None,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ndarray[tuple[int], np.dtype[np.void]]: ...
|
|
@overload
|
|
def join_by(
|
|
key: str | Sequence[str],
|
|
r1: npt.NDArray[np.void],
|
|
r2: npt.NDArray[np.void],
|
|
jointype: _JoinType = "inner",
|
|
r1postfix: str = "1",
|
|
r2postfix: str = "2",
|
|
defaults: Mapping[str, object] | None = None,
|
|
*,
|
|
usemask: Literal[False],
|
|
asrecarray: Literal[True],
|
|
) -> np.recarray[tuple[int], np.dtype[np.void]]: ...
|
|
@overload
|
|
def join_by(
|
|
key: str | Sequence[str],
|
|
r1: npt.NDArray[np.void],
|
|
r2: npt.NDArray[np.void],
|
|
jointype: _JoinType = "inner",
|
|
r1postfix: str = "1",
|
|
r2postfix: str = "2",
|
|
defaults: Mapping[str, object] | None = None,
|
|
usemask: Literal[True] = True,
|
|
asrecarray: Literal[False] = False,
|
|
) -> np.ma.MaskedArray[tuple[int], np.dtype[np.void]]: ...
|
|
@overload
|
|
def join_by(
|
|
key: str | Sequence[str],
|
|
r1: npt.NDArray[np.void],
|
|
r2: npt.NDArray[np.void],
|
|
jointype: _JoinType = "inner",
|
|
r1postfix: str = "1",
|
|
r2postfix: str = "2",
|
|
defaults: Mapping[str, object] | None = None,
|
|
usemask: Literal[True] = True,
|
|
*,
|
|
asrecarray: Literal[True],
|
|
) -> MaskedRecords[tuple[int], np.dtype[np.void]]: ...
|
|
|
|
#
|
|
def rec_join(
|
|
key: str | Sequence[str],
|
|
r1: npt.NDArray[np.void],
|
|
r2: npt.NDArray[np.void],
|
|
jointype: _JoinType = "inner",
|
|
r1postfix: str = "1",
|
|
r2postfix: str = "2",
|
|
defaults: Mapping[str, object] | None = None,
|
|
) -> np.recarray[tuple[int], np.dtype[np.void]]: ...
|