Repository: numpy/numpy
File: numpy/lib/tests/test_function_base.py (Python)

Code before:
        # TODO: Note that times have dubious rounding as of fixing NaTs!
        # NaT and NaN should behave the same, do basic tests for NaT:
        a = np.arange(0, 24, dtype=dtype)
        a[pos] = "NaT"
        res = np.percentile(a, 30)
        assert res.dtype == dtype
        assert np.isnat(res)
        res = np.percentile(a, [30, 60])
        assert res.dtype == dtype
        assert np.isnat(res).all()

        a = np.arange(0, 24 * 3, dtype=dtype).reshape(-1, 3)
        a[pos, 1] = "NaT"
        res = np.percentile(a, 30, axis=0)
        assert_array_equal(np.isnat(res), [False, True, False])

    @pytest.mark.parametrize("qtype", [np.float16, np.float32])
    @pytest.mark.parametrize("method", quantile_methods)
    def test_percentile_gh_29003(self, qtype, method):
        zero = qtype(0)
        one = qtype(0)
        data = [zero] * 65521
        a = np.array(data)
        a[:20_000] = one
        z = np.percentile(a, 50, method=method)
        assert z == zero
        assert z.dtype == a.dtype
        z = np.percentile(a, .9, method=method)
        assert z == one
        assert z.dtype == a.dtype

    def test_percentile_gh_29003_Fraction(self):
        zero = Fraction(0)
        one = Fraction(0)
        data = [zero] * 65521
        a = np.array(data)
        a[:20_000] = one
        z = np.percentile(a, 50)
        assert z == zero
        z = np.percentile(a, Fraction(50))
        assert z == zero
        assert np.array(z).dtype == a.dtype


class TestQuantile:
    # most of this is already tested by TestPercentile

    def V(self, x, y, alpha):
        # Identification function used in several tests.
        return (x >= y) - alpha


Review comment (suggestion):
Just `z == 0` -- there is no benefit to using the correct type in the assert statement.

Code after:
        # TODO: Note that times have dubious rounding as of fixing NaTs!
        # NaT and NaN should behave the same, do basic tests for NaT:
        a = np.arange(0, 24, dtype=dtype)
        a[pos] = "NaT"
        res = np.percentile(a, 30)
        assert res.dtype == dtype
        assert np.isnat(res)
        res = np.percentile(a, [30, 60])
        assert res.dtype == dtype
        assert np.isnat(res).all()

        a = np.arange(0, 24 * 3, dtype=dtype).reshape(-1, 3)
        a[pos, 1] = "NaT"
        res = np.percentile(a, 30, axis=0)
        assert_array_equal(np.isnat(res), [False, True, False])

    @pytest.mark.parametrize("qtype", [np.float16, np.float32])
    @pytest.mark.parametrize("method", quantile_methods)
    def test_percentile_gh_29003(self, qtype, method):
        # test that with float16 or float32 input we do not get overflow
        zero = qtype(0)
        one = qtype(1)
        a = np.zeros(65521, qtype)
        a[:20_000] = one
        z = np.percentile(a, 50, method=method)
        assert z == zero
        assert z.dtype == a.dtype
        z = np.percentile(a, 99, method=method)
        assert z == one
        assert z.dtype == a.dtype

    def test_percentile_gh_29003_Fraction(self):
        zero = Fraction(0)
        one = Fraction(1)
        a = np.array([zero] * 65521)
        a[:20_000] = one
        z = np.percentile(a, 50)
        assert z == zero
        z = np.percentile(a, Fraction(50))
        assert z == zero
        assert np.array(z).dtype == a.dtype

        z = np.percentile(a, 99)
        assert z == one
        # test that with only Fraction input the return type is a Fraction
        z = np.percentile(a, Fraction(99))
        assert z == one
        assert np.array(z).dtype == a.dtype


class TestQuantile: