Repository: numpy/numpy
File: numpy/typing/tests/data/reveal/fromnumeric.pyi (Python)

Code before:
assert_type(np.argpartition(b, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(f4, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(f, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(AR_b, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(AR_f4, 0), npt.NDArray[np.intp])

assert_type(np.sort([2, 1], 0), npt.NDArray[Any])
assert_type(np.sort(AR_b, 0), npt.NDArray[np.bool])
assert_type(np.sort(AR_f4, 0), npt.NDArray[np.float32])

assert_type(np.argsort(AR_b, 0), npt.NDArray[np.intp])
assert_type(np.argsort(AR_f4, 0), npt.NDArray[np.intp])

assert_type(np.argmax(AR_b), np.intp)
assert_type(np.argmax(AR_f4), np.intp)
assert_type(np.argmax(AR_b, axis=0), Any)
assert_type(np.argmax(AR_f4, axis=0), Any)
assert_type(np.argmax(AR_f4, out=AR_subclass), NDArraySubclass)

assert_type(np.argmin(AR_b), np.intp)
assert_type(np.argmin(AR_f4), np.intp)
assert_type(np.argmin(AR_b, axis=0), Any)
assert_type(np.argmin(AR_f4, axis=0), Any)
assert_type(np.argmin(AR_f4, out=AR_subclass), NDArraySubclass)

OUT_I = np.empty((), dtype=np.intp)
OUT_F = np.empty((), dtype=np.float64)
assert_type(np.argmin(AR_f4, out=OUT_I), npt.NDArray[np.intp])
assert_type(np.argmax(AR_f4, out=OUT_I), npt.NDArray[np.intp])

assert_type(np.searchsorted(AR_b[0], 0), np.intp)
assert_type(np.searchsorted(AR_f4[0], 0), np.intp)
assert_type(np.searchsorted(AR_b[0], [0]), npt.NDArray[np.intp])
assert_type(np.searchsorted(AR_f4[0], [0]), npt.NDArray[np.intp])

assert_type(np.resize(b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.resize(f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.resize(f, (5, 5)), np.ndarray[tuple[int, int], np.dtype])
assert_type(np.resize(AR_b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.resize(AR_f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])

assert_type(np.squeeze(b), np.bool)
assert_type(np.squeeze(f4), np.float32)
assert_type(np.squeeze(f), npt.NDArray[Any])
assert_type(np.squeeze(AR_b), npt.NDArray[np.bool])
assert_type(np.squeeze(AR_f4), npt.NDArray[np.float32])

assert_type(np.diagonal(AR_b), npt.NDArray[np.bool])
assert_type(np.diagonal(AR_f4), npt.NDArray[np.float32])

assert_type(np.trace(AR_b), Any)

Review comment (bug):
There's some shape-typing involved here, which is causing the mypy tests to fail. It's probably easiest to just use `AR_i8` instead.

Code after:
assert_type(np.partition(f, 0, axis=None), npt.NDArray[Any])
assert_type(np.partition(AR_b, 0), npt.NDArray[np.bool])
assert_type(np.partition(AR_f4, 0), npt.NDArray[np.float32])

assert_type(np.argpartition(b, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(f4, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(f, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(AR_b, 0), npt.NDArray[np.intp])
assert_type(np.argpartition(AR_f4, 0), npt.NDArray[np.intp])

assert_type(np.sort([2, 1], 0), npt.NDArray[Any])
assert_type(np.sort(AR_b, 0), npt.NDArray[np.bool])
assert_type(np.sort(AR_f4, 0), npt.NDArray[np.float32])

assert_type(np.argsort(AR_b, 0), npt.NDArray[np.intp])
assert_type(np.argsort(AR_f4, 0), npt.NDArray[np.intp])

assert_type(np.argmax(AR_b), np.intp)
assert_type(np.argmax(AR_f4), np.intp)
assert_type(np.argmax(AR_b, axis=0), Any)
assert_type(np.argmax(AR_f4, axis=0), Any)
assert_type(np.argmax(AR_f4, out=AR_sub_i), NDArrayIntSubclass)

assert_type(np.argmin(AR_b), np.intp)
assert_type(np.argmin(AR_f4), np.intp)
assert_type(np.argmin(AR_b, axis=0), Any)
assert_type(np.argmin(AR_f4, axis=0), Any)
assert_type(np.argmin(AR_f4, out=AR_sub_i), NDArrayIntSubclass)

assert_type(np.searchsorted(AR_b[0], 0), np.intp)
assert_type(np.searchsorted(AR_f4[0], 0), np.intp)
assert_type(np.searchsorted(AR_b[0], [0]), npt.NDArray[np.intp])
assert_type(np.searchsorted(AR_f4[0], [0]), npt.NDArray[np.intp])

assert_type(np.resize(b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.resize(f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.resize(f, (5, 5)), np.ndarray[tuple[int, int], np.dtype])
assert_type(np.resize(AR_b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
assert_type(np.resize(AR_f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])

assert_type(np.squeeze(b), np.bool)
assert_type(np.squeeze(f4), np.float32)
assert_type(np.squeeze(f), npt.NDArray[Any])
assert_type(np.squeeze(AR_b), npt.NDArray[np.bool])
assert_type(np.squeeze(AR_f4), npt.NDArray[np.float32])

assert_type(np.diagonal(AR_b), npt.NDArray[np.bool])
assert_type(np.diagonal(AR_f4), npt.NDArray[np.float32])

assert_type(np.trace(AR_b), Any)
assert_type(np.trace(AR_f4), Any)