arraysetops.pyi 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. from typing import Any
  2. import numpy as np
  3. import numpy.typing as npt
  4. from numpy.lib._arraysetops_impl import UniqueAllResult, UniqueCountsResult, UniqueInverseResult
  5. from typing_extensions import assert_type
  6. AR_b: npt.NDArray[np.bool]
  7. AR_i8: npt.NDArray[np.int64]
  8. AR_f8: npt.NDArray[np.float64]
  9. AR_M: npt.NDArray[np.datetime64]
  10. AR_O: npt.NDArray[np.object_]
  11. AR_LIKE_f8: list[float]
  12. assert_type(np.ediff1d(AR_b), npt.NDArray[np.int8])
  13. assert_type(np.ediff1d(AR_i8, to_end=[1, 2, 3]), npt.NDArray[np.int64])
  14. assert_type(np.ediff1d(AR_M), npt.NDArray[np.timedelta64])
  15. assert_type(np.ediff1d(AR_O), npt.NDArray[np.object_])
  16. assert_type(np.ediff1d(AR_LIKE_f8, to_begin=[1, 1.5]), npt.NDArray[Any])
  17. assert_type(np.intersect1d(AR_i8, AR_i8), npt.NDArray[np.int64])
  18. assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64])
  19. assert_type(np.intersect1d(AR_f8, AR_i8), npt.NDArray[Any])
  20. assert_type(
  21. np.intersect1d(AR_f8, AR_f8, return_indices=True),
  22. tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]],
  23. )
  24. assert_type(np.setxor1d(AR_i8, AR_i8), npt.NDArray[np.int64])
  25. assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64])
  26. assert_type(np.setxor1d(AR_f8, AR_i8), npt.NDArray[Any])
  27. assert_type(np.isin(AR_i8, AR_i8), npt.NDArray[np.bool])
  28. assert_type(np.isin(AR_M, AR_M, assume_unique=True), npt.NDArray[np.bool])
  29. assert_type(np.isin(AR_f8, AR_i8), npt.NDArray[np.bool])
  30. assert_type(np.isin(AR_f8, AR_LIKE_f8, invert=True), npt.NDArray[np.bool])
  31. assert_type(np.union1d(AR_i8, AR_i8), npt.NDArray[np.int64])
  32. assert_type(np.union1d(AR_M, AR_M), npt.NDArray[np.datetime64])
  33. assert_type(np.union1d(AR_f8, AR_i8), npt.NDArray[Any])
  34. assert_type(np.setdiff1d(AR_i8, AR_i8), npt.NDArray[np.int64])
  35. assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64])
  36. assert_type(np.setdiff1d(AR_f8, AR_i8), npt.NDArray[Any])
  37. assert_type(np.unique(AR_f8), npt.NDArray[np.float64])
  38. assert_type(np.unique(AR_LIKE_f8, axis=0), npt.NDArray[Any])
  39. assert_type(np.unique(AR_f8, return_index=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
  40. assert_type(np.unique(AR_LIKE_f8, return_index=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
  41. assert_type(np.unique(AR_f8, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
  42. assert_type(np.unique(AR_LIKE_f8, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
  43. assert_type(np.unique(AR_f8, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
  44. assert_type(np.unique(AR_LIKE_f8, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
  45. assert_type(np.unique(AR_f8, return_index=True, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  46. assert_type(np.unique(AR_LIKE_f8, return_index=True, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  47. assert_type(np.unique(AR_f8, return_index=True, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  48. assert_type(np.unique(AR_LIKE_f8, return_index=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  49. assert_type(np.unique(AR_f8, return_inverse=True, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  50. assert_type(np.unique(AR_LIKE_f8, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  51. assert_type(np.unique(AR_f8, return_index=True, return_inverse=True, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  52. assert_type(np.unique(AR_LIKE_f8, return_index=True, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]])
  53. assert_type(np.unique_all(AR_f8), UniqueAllResult[np.float64])
  54. assert_type(np.unique_all(AR_LIKE_f8), UniqueAllResult[Any])
  55. assert_type(np.unique_counts(AR_f8), UniqueCountsResult[np.float64])
  56. assert_type(np.unique_counts(AR_LIKE_f8), UniqueCountsResult[Any])
  57. assert_type(np.unique_inverse(AR_f8), UniqueInverseResult[np.float64])
  58. assert_type(np.unique_inverse(AR_LIKE_f8), UniqueInverseResult[Any])
  59. assert_type(np.unique_values(AR_f8), npt.NDArray[np.float64])
  60. assert_type(np.unique_values(AR_LIKE_f8), npt.NDArray[Any])