Repository: numpy/numpy
File: numpy/_core/tests/test_mem_overlap.py (Python)

Code before:
    assert_equal(np.shares_memory(x, x.copy()), False)

    a = x[:, ::2, ::3]
    b = x[:, ::3, ::2]
    assert_equal(np.shares_memory(a, b), True)
    assert_equal(np.shares_memory(a, b, max_work=None), True)
    assert_raises(
        np.exceptions.TooHardError, np.shares_memory, a, b, max_work=1
    )


def test_may_share_memory_bad_max_work():
    x = np.zeros([1])
    assert_raises(OverflowError, np.may_share_memory, x, x, max_work=10**100)
    assert_raises(OverflowError, np.shares_memory, x, x, max_work=10**100)


def test_internal_overlap_diophantine():
    def check(A, U, exists=None):
        X = solve_diophantine(A, U, 0, require_ub_nontrivial=1)

        if exists is None:
            exists = (X is not None)

        if X is not None:
            assert_(sum(a * x for a, x in zip(A, X)) == sum(a * u // 2 for a, u in zip(A, U)))  # noqa: E501
            assert_(all(0 <= x <= u for x, u in zip(X, U)))
            assert_(any(x != u // 2 for x, u in zip(X, U)))

        if exists:
            assert_(X is not None, repr(X))
        else:
            assert_(X is None, repr(X))

    # Smoke tests
    check((3, 2), (2 * 2, 3 * 2), exists=True)
    check((3 * 2, 2), (15 * 2, (3 - 1) * 2), exists=False)


def test_internal_overlap_slices():
    # Slicing an array never generates internal overlap

    x = np.zeros([17, 34, 71, 97], dtype=np.int16)

    rng = np.random.RandomState(1234)

    def random_slice(n, step):
        start = rng.randint(0, n + 1, dtype=np.intp)
        stop = rng.randint(start, n + 1, dtype=np.intp)
        if rng.randint(0, 2, dtype=np.intp) == 0:
            stop, start = start, stop

Review comment (suggestion):
I'd break this into two computations before the `assert_`. Now that we are using pytest, you can also use `assert` instead of the `assert_` function.

Code after:
    assert_equal(np.shares_memory(x, x), True)
    assert_equal(np.shares_memory(x, x.copy()), False)

    a = x[:, ::2, ::3]
    b = x[:, ::3, ::2]
    assert_equal(np.shares_memory(a, b), True)
    assert_equal(np.shares_memory(a, b, max_work=None), True)
    assert_raises(
        np.exceptions.TooHardError, np.shares_memory, a, b, max_work=1
    )


def test_may_share_memory_bad_max_work():
    x = np.zeros([1])
    assert_raises(OverflowError, np.may_share_memory, x, x, max_work=10**100)
    assert_raises(OverflowError, np.shares_memory, x, x, max_work=10**100)


def test_internal_overlap_diophantine():
    def check(A, U, exists=None):
        X = solve_diophantine(A, U, 0, require_ub_nontrivial=1)

        if exists is None:
            exists = (X is not None)

        if X is not None:
            sum_ax = sum(a * x for a, x in zip(A, X))
            sum_au_half = sum(a * u // 2 for a, u in zip(A, U))
            assert_(sum_ax == sum_au_half)
            assert_(all(0 <= x <= u for x, u in zip(X, U)))
            assert_(any(x != u // 2 for x, u in zip(X, U)))

        if exists:
            assert_(X is not None, repr(X))
        else:
            assert_(X is None, repr(X))

    # Smoke tests
    check((3, 2), (2 * 2, 3 * 2), exists=True)
    check((3 * 2, 2), (15 * 2, (3 - 1) * 2), exists=False)


def test_internal_overlap_slices():
    # Slicing an array never generates internal overlap

    x = np.zeros([17, 34, 71, 97], dtype=np.int16)

    rng = np.random.RandomState(1234)

    def random_slice(n, step):
        start = rng.randint(0, n + 1, dtype=np.intp)