You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
109 lines
2.7 KiB
109 lines
2.7 KiB
1 year ago
|
"""
|
||
|
Tests that work on both the Python and C engines but do not have a
|
||
|
specific classification into the other test modules.
|
||
|
"""
|
||
|
from io import StringIO
|
||
|
|
||
|
import pytest
|
||
|
|
||
|
from pandas import (
|
||
|
DataFrame,
|
||
|
concat,
|
||
|
)
|
||
|
import pandas._testing as tm
|
||
|
|
||
|
pytestmark = pytest.mark.usefixtures("pyarrow_skip")
|
||
|
|
||
|
|
||
|
def test_iterator(all_parsers):
|
||
|
# see gh-6607
|
||
|
data = """index,A,B,C,D
|
||
|
foo,2,3,4,5
|
||
|
bar,7,8,9,10
|
||
|
baz,12,13,14,15
|
||
|
qux,12,13,14,15
|
||
|
foo2,12,13,14,15
|
||
|
bar2,12,13,14,15
|
||
|
"""
|
||
|
parser = all_parsers
|
||
|
kwargs = {"index_col": 0}
|
||
|
|
||
|
expected = parser.read_csv(StringIO(data), **kwargs)
|
||
|
with parser.read_csv(StringIO(data), iterator=True, **kwargs) as reader:
|
||
|
first_chunk = reader.read(3)
|
||
|
tm.assert_frame_equal(first_chunk, expected[:3])
|
||
|
|
||
|
last_chunk = reader.read(5)
|
||
|
tm.assert_frame_equal(last_chunk, expected[3:])
|
||
|
|
||
|
|
||
|
def test_iterator2(all_parsers):
|
||
|
parser = all_parsers
|
||
|
data = """A,B,C
|
||
|
foo,1,2,3
|
||
|
bar,4,5,6
|
||
|
baz,7,8,9
|
||
|
"""
|
||
|
|
||
|
with parser.read_csv(StringIO(data), iterator=True) as reader:
|
||
|
result = list(reader)
|
||
|
|
||
|
expected = DataFrame(
|
||
|
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
|
||
|
index=["foo", "bar", "baz"],
|
||
|
columns=["A", "B", "C"],
|
||
|
)
|
||
|
tm.assert_frame_equal(result[0], expected)
|
||
|
|
||
|
|
||
|
def test_iterator_stop_on_chunksize(all_parsers):
|
||
|
# gh-3967: stopping iteration when chunksize is specified
|
||
|
parser = all_parsers
|
||
|
data = """A,B,C
|
||
|
foo,1,2,3
|
||
|
bar,4,5,6
|
||
|
baz,7,8,9
|
||
|
"""
|
||
|
|
||
|
with parser.read_csv(StringIO(data), chunksize=1) as reader:
|
||
|
result = list(reader)
|
||
|
|
||
|
assert len(result) == 3
|
||
|
expected = DataFrame(
|
||
|
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
|
||
|
index=["foo", "bar", "baz"],
|
||
|
columns=["A", "B", "C"],
|
||
|
)
|
||
|
tm.assert_frame_equal(concat(result), expected)
|
||
|
|
||
|
|
||
|
@pytest.mark.parametrize(
|
||
|
"kwargs", [{"iterator": True, "chunksize": 1}, {"iterator": True}, {"chunksize": 1}]
|
||
|
)
|
||
|
def test_iterator_skipfooter_errors(all_parsers, kwargs):
|
||
|
msg = "'skipfooter' not supported for iteration"
|
||
|
parser = all_parsers
|
||
|
data = "a\n1\n2"
|
||
|
|
||
|
with pytest.raises(ValueError, match=msg):
|
||
|
with parser.read_csv(StringIO(data), skipfooter=1, **kwargs) as _:
|
||
|
pass
|
||
|
|
||
|
|
||
|
def test_iteration_open_handle(all_parsers):
|
||
|
parser = all_parsers
|
||
|
kwargs = {"header": None}
|
||
|
|
||
|
with tm.ensure_clean() as path:
|
||
|
with open(path, "w", encoding="utf-8") as f:
|
||
|
f.write("AAA\nBBB\nCCC\nDDD\nEEE\nFFF\nGGG")
|
||
|
|
||
|
with open(path, encoding="utf-8") as f:
|
||
|
for line in f:
|
||
|
if "CCC" in line:
|
||
|
break
|
||
|
|
||
|
result = parser.read_csv(f, **kwargs)
|
||
|
expected = DataFrame({0: ["DDD", "EEE", "FFF", "GGG"]})
|
||
|
tm.assert_frame_equal(result, expected)
|