| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162 |
- """These tests are based on the doctests from `numpy/lib/recfunctions.py`."""
- from typing import Any
- from typing_extensions import assert_type
- import numpy as np
- import numpy.typing as npt
- from numpy.lib import recfunctions as rfn
- def test_recursive_fill_fields() -> None:
- a: npt.NDArray[np.void] = np.array(
- [(1, 10.0), (2, 20.0)],
- dtype=[("A", np.int64), ("B", np.float64)],
- )
- b = np.zeros((int(3),), dtype=a.dtype)
- out = rfn.recursive_fill_fields(a, b)
- assert_type(out, np.ndarray[tuple[int], np.dtype[np.void]])
- def test_get_names() -> None:
- names: tuple[str | Any, ...]
- names = rfn.get_names(np.empty((1,), dtype=[("A", int)]).dtype)
- names = rfn.get_names(np.empty((1,), dtype=[("A", int), ("B", float)]).dtype)
- adtype = np.dtype([("a", int), ("b", [("b_a", int), ("b_b", int)])])
- names = rfn.get_names(adtype)
- def test_get_names_flat() -> None:
- names: tuple[str, ...]
- names = rfn.get_names_flat(np.empty((1,), dtype=[("A", int)]).dtype)
- names = rfn.get_names_flat(np.empty((1,), dtype=[("A", int), ("B", float)]).dtype)
- adtype = np.dtype([("a", int), ("b", [("b_a", int), ("b_b", int)])])
- names = rfn.get_names_flat(adtype)
- def test_flatten_descr() -> None:
- ndtype = np.dtype([("a", "<i4"), ("b", [("b_a", "<f8"), ("b_b", "<i4")])])
- assert_type(rfn.flatten_descr(ndtype), tuple[tuple[str, np.dtype[Any]]])
- def test_get_fieldstructure() -> None:
- ndtype = np.dtype([
- ("A", int),
- ("B", [("B_A", int), ("B_B", [("B_B_A", int), ("B_B_B", int)])]),
- ])
- assert_type(rfn.get_fieldstructure(ndtype), dict[str, list[str]])
- def test_merge_arrays() -> None:
- assert_type(
- rfn.merge_arrays((
- np.ones((int(2),), np.int_),
- np.ones((int(3),), np.float64),
- )),
- np.recarray[tuple[int], np.dtype[np.void]],
- )
- def test_drop_fields() -> None:
- ndtype = [("a", np.int64), ("b", [("b_a", np.double), ("b_b", np.int64)])]
- a = np.ones((int(3),), dtype=ndtype)
- assert_type(
- rfn.drop_fields(a, "a"),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- assert_type(
- rfn.drop_fields(a, "a", asrecarray=True),
- np.rec.recarray[tuple[int], np.dtype[np.void]],
- )
- assert_type(
- rfn.rec_drop_fields(a, "a"),
- np.rec.recarray[tuple[int], np.dtype[np.void]],
- )
- def test_rename_fields() -> None:
- ndtype = [("a", np.int64), ("b", [("b_a", np.double), ("b_b", np.int64)])]
- a = np.ones((int(3),), dtype=ndtype)
- assert_type(
- rfn.rename_fields(a, {"a": "A", "b_b": "B_B"}),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- def test_repack_fields() -> None:
- dt: np.dtype[np.void] = np.dtype("u1, <i8, <f8", align=True)
- assert_type(rfn.repack_fields(dt), np.dtype[np.void])
- assert_type(rfn.repack_fields(dt.type(0)), np.void)
- assert_type(
- rfn.repack_fields(np.ones((int(3),), dtype=dt)),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- def test_structured_to_unstructured() -> None:
- a = np.zeros(4, dtype=[("a", "i4"), ("b", "f4,u2"), ("c", "f4", 2)])
- assert_type(rfn.structured_to_unstructured(a), npt.NDArray[Any])
- def unstructured_to_structured() -> None:
- dt: np.dtype[np.void] = np.dtype([("a", "i4"), ("b", "f4,u2"), ("c", "f4", 2)])
- a = np.arange(20, dtype=np.int32).reshape((4, 5))
- assert_type(rfn.unstructured_to_structured(a, dt), npt.NDArray[np.void])
- def test_apply_along_fields() -> None:
- b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
- assert_type(
- rfn.apply_along_fields(np.mean, b),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- def test_assign_fields_by_name() -> None:
- b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
- assert_type(
- rfn.apply_along_fields(np.mean, b),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- def test_require_fields() -> None:
- a = np.ones(4, dtype=[("a", "i4"), ("b", "f8"), ("c", "u1")])
- assert_type(
- rfn.require_fields(a, [("b", "f4"), ("c", "u1")]),
- np.ndarray[tuple[int], np.dtype[np.void]],
- )
- def test_stack_arrays() -> None:
- x = np.zeros((int(2),), np.int32)
- assert_type(
- rfn.stack_arrays(x),
- np.ndarray[tuple[int], np.dtype[np.int32]],
- )
- z = np.ones((int(2),), [("A", "|S3"), ("B", float)])
- zz = np.ones((int(2),), [("A", "|S3"), ("B", np.float64), ("C", np.float64)])
- assert_type(
- rfn.stack_arrays((z, zz)),
- np.ma.MaskedArray[tuple[int, ...], np.dtype[np.void]],
- )
- def test_find_duplicates() -> None:
- ndtype = np.dtype([("a", int)])
- a = np.ma.ones(7, mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype)
- assert_type(rfn.find_duplicates(a), np.ma.MaskedArray[Any, np.dtype[np.void]])
- assert_type(
- rfn.find_duplicates(a, ignoremask=True, return_index=True),
- tuple[
- np.ma.MaskedArray[Any, np.dtype[np.void]],
- np.ndarray[Any, np.dtype[np.int_]],
- ],
- )
|