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

Code before:
    res = np.loadtxt(txt, dtype=str, delimiter=" ", max_rows=nmax)
    assert len(res) == nmax

    # file-obj path
    fd, fname = mkstemp()
    os.close(fd)
    with open(fname, "w") as fh:
        fh.write("\n".join(data))
    res = np.loadtxt(fname, dtype=str, delimiter=" ", max_rows=nmax)
    os.remove(fname)
    assert len(res) == nmax

@pytest.mark.parametrize("nskip", (0, 10000, 50000, 100000))
def test_skiprow_exceeding_maxrows_exceeding_chunksize(tmpdir, nskip):
    # tries to read all of the file,
    # or less, equal, greater than _loadtxt_chunksize
    file_length = 110000
    data = np.arange(1, file_length + 1).astype(str) + np.array([" a 0.5 1"] * file_length)

    # file-like path
    txt = StringIO("\n".join(data))
    res = np.loadtxt(txt, dtype = 'str', delimiter=" ", skiprows=nskip, max_rows=60000)
    expected_length = min(60000, file_length - nskip)
    assert len(res) == expected_length
    # are the right lines read in res?
    assert np.all(np.equal(np.arange(nskip + 1, nskip + 1 + expected_length).astype(str), res[:, 0]))

    # file-obj path
    tmp_file = tmpdir.join("test_data.txt")
    tmp_file.write("\n".join(data))
    fname = str(tmp_file)
    res = np.loadtxt(fname, dtype = 'str', delimiter=" ", skiprows=nskip, max_rows=60000)
    expected_length = min(60000, file_length - nskip)
    assert len(res) == expected_length
    # are the right lines read in res?
    assert np.all(np.equal(np.arange(nskip + 1, nskip + 1 + expected_length).astype(str), res[:, 0]))


Review comment (suggestion):
Just use `assert_array_equal(expected, res[:, 0])`.  And set `expected` earlier so we can re-use it.  (I think it is still fine to explicitly check the length, since that is the main thing)

Code after:
    res = np.loadtxt(txt, dtype=str, delimiter=" ", max_rows=nmax)
    assert len(res) == nmax

    # file-obj path
    fd, fname = mkstemp()
    os.close(fd)
    with open(fname, "w") as fh:
        fh.write("\n".join(data))
    res = np.loadtxt(fname, dtype=str, delimiter=" ", max_rows=nmax)
    os.remove(fname)
    assert len(res) == nmax

@pytest.mark.parametrize("nskip", (0, 10000, 12345, 50000, 67891, 100000))
def test_skiprow_exceeding_maxrows_exceeding_chunksize(tmpdir, nskip):
    # tries to read a file in chunks by skipping a variable amount of lines,
    # less, equal, greater than max_rows
    file_length = 110000
    data = "\n".join(f"{i} a 0.5 1" for i in range(1, file_length + 1))
    expected_length = min(60000, file_length - nskip)
    expected = np.arange(nskip + 1, nskip + 1 + expected_length).astype(str)

    # file-like path
    txt = StringIO(data)
    res = np.loadtxt(txt, dtype='str', delimiter=" ", skiprows=nskip, max_rows=60000)
    assert len(res) == expected_length
    # are the right lines read in res?
    assert_array_equal(expected, res[:, 0])

    # file-obj path
    tmp_file = tmpdir / "test_data.txt"
    tmp_file.write(data)
    fname = str(tmp_file)
    res = np.loadtxt(fname, dtype='str', delimiter=" ", skiprows=nskip, max_rows=60000)
    assert len(res) == expected_length
    # are the right lines read in res?
    assert_array_equal(expected, res[:, 0])