recfunctions.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. """These tests are based on the doctests from `numpy/lib/recfunctions.py`."""
  2. from typing import Any
  3. from typing_extensions import assert_type
  4. import numpy as np
  5. import numpy.typing as npt
  6. from numpy.lib import recfunctions as rfn
  7. def test_recursive_fill_fields() -> None:
  8. a: npt.NDArray[np.void] = np.array(
  9. [(1, 10.0), (2, 20.0)],
  10. dtype=[("A", np.int64), ("B", np.float64)],
  11. )
  12. b = np.zeros((int(3),), dtype=a.dtype)
  13. out = rfn.recursive_fill_fields(a, b)
  14. assert_type(out, np.ndarray[tuple[int], np.dtype[np.void]])
  15. def test_get_names() -> None:
  16. names: tuple[str | Any, ...]
  17. names = rfn.get_names(np.empty((1,), dtype=[("A", int)]).dtype)
  18. names = rfn.get_names(np.empty((1,), dtype=[("A", int), ("B", float)]).dtype)
  19. adtype = np.dtype([("a", int), ("b", [("b_a", int), ("b_b", int)])])
  20. names = rfn.get_names(adtype)
  21. def test_get_names_flat() -> None:
  22. names: tuple[str, ...]
  23. names = rfn.get_names_flat(np.empty((1,), dtype=[("A", int)]).dtype)
  24. names = rfn.get_names_flat(np.empty((1,), dtype=[("A", int), ("B", float)]).dtype)
  25. adtype = np.dtype([("a", int), ("b", [("b_a", int), ("b_b", int)])])
  26. names = rfn.get_names_flat(adtype)
  27. def test_flatten_descr() -> None:
  28. ndtype = np.dtype([("a", "<i4"), ("b", [("b_a", "<f8"), ("b_b", "<i4")])])
  29. assert_type(rfn.flatten_descr(ndtype), tuple[tuple[str, np.dtype[Any]]])
  30. def test_get_fieldstructure() -> None:
  31. ndtype = np.dtype([
  32. ("A", int),
  33. ("B", [("B_A", int), ("B_B", [("B_B_A", int), ("B_B_B", int)])]),
  34. ])
  35. assert_type(rfn.get_fieldstructure(ndtype), dict[str, list[str]])
  36. def test_merge_arrays() -> None:
  37. assert_type(
  38. rfn.merge_arrays((
  39. np.ones((int(2),), np.int_),
  40. np.ones((int(3),), np.float64),
  41. )),
  42. np.recarray[tuple[int], np.dtype[np.void]],
  43. )
  44. def test_drop_fields() -> None:
  45. ndtype = [("a", np.int64), ("b", [("b_a", np.double), ("b_b", np.int64)])]
  46. a = np.ones((int(3),), dtype=ndtype)
  47. assert_type(
  48. rfn.drop_fields(a, "a"),
  49. np.ndarray[tuple[int], np.dtype[np.void]],
  50. )
  51. assert_type(
  52. rfn.drop_fields(a, "a", asrecarray=True),
  53. np.rec.recarray[tuple[int], np.dtype[np.void]],
  54. )
  55. assert_type(
  56. rfn.rec_drop_fields(a, "a"),
  57. np.rec.recarray[tuple[int], np.dtype[np.void]],
  58. )
  59. def test_rename_fields() -> None:
  60. ndtype = [("a", np.int64), ("b", [("b_a", np.double), ("b_b", np.int64)])]
  61. a = np.ones((int(3),), dtype=ndtype)
  62. assert_type(
  63. rfn.rename_fields(a, {"a": "A", "b_b": "B_B"}),
  64. np.ndarray[tuple[int], np.dtype[np.void]],
  65. )
  66. def test_repack_fields() -> None:
  67. dt: np.dtype[np.void] = np.dtype("u1, <i8, <f8", align=True)
  68. assert_type(rfn.repack_fields(dt), np.dtype[np.void])
  69. assert_type(rfn.repack_fields(dt.type(0)), np.void)
  70. assert_type(
  71. rfn.repack_fields(np.ones((int(3),), dtype=dt)),
  72. np.ndarray[tuple[int], np.dtype[np.void]],
  73. )
  74. def test_structured_to_unstructured() -> None:
  75. a = np.zeros(4, dtype=[("a", "i4"), ("b", "f4,u2"), ("c", "f4", 2)])
  76. assert_type(rfn.structured_to_unstructured(a), npt.NDArray[Any])
  77. def unstructured_to_structured() -> None:
  78. dt: np.dtype[np.void] = np.dtype([("a", "i4"), ("b", "f4,u2"), ("c", "f4", 2)])
  79. a = np.arange(20, dtype=np.int32).reshape((4, 5))
  80. assert_type(rfn.unstructured_to_structured(a, dt), npt.NDArray[np.void])
  81. def test_apply_along_fields() -> None:
  82. b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
  83. assert_type(
  84. rfn.apply_along_fields(np.mean, b),
  85. np.ndarray[tuple[int], np.dtype[np.void]],
  86. )
  87. def test_assign_fields_by_name() -> None:
  88. b = np.ones(4, dtype=[("x", "i4"), ("y", "f4"), ("z", "f8")])
  89. assert_type(
  90. rfn.apply_along_fields(np.mean, b),
  91. np.ndarray[tuple[int], np.dtype[np.void]],
  92. )
  93. def test_require_fields() -> None:
  94. a = np.ones(4, dtype=[("a", "i4"), ("b", "f8"), ("c", "u1")])
  95. assert_type(
  96. rfn.require_fields(a, [("b", "f4"), ("c", "u1")]),
  97. np.ndarray[tuple[int], np.dtype[np.void]],
  98. )
  99. def test_stack_arrays() -> None:
  100. x = np.zeros((int(2),), np.int32)
  101. assert_type(
  102. rfn.stack_arrays(x),
  103. np.ndarray[tuple[int], np.dtype[np.int32]],
  104. )
  105. z = np.ones((int(2),), [("A", "|S3"), ("B", float)])
  106. zz = np.ones((int(2),), [("A", "|S3"), ("B", np.float64), ("C", np.float64)])
  107. assert_type(
  108. rfn.stack_arrays((z, zz)),
  109. np.ma.MaskedArray[tuple[int, ...], np.dtype[np.void]],
  110. )
  111. def test_find_duplicates() -> None:
  112. ndtype = np.dtype([("a", int)])
  113. a = np.ma.ones(7, mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype)
  114. assert_type(rfn.find_duplicates(a), np.ma.MaskedArray[Any, np.dtype[np.void]])
  115. assert_type(
  116. rfn.find_duplicates(a, ignoremask=True, return_index=True),
  117. tuple[
  118. np.ma.MaskedArray[Any, np.dtype[np.void]],
  119. np.ndarray[Any, np.dtype[np.int_]],
  120. ],
  121. )