Files
RustPython/extra_tests/snippets/stdlib_csv.py
T
2026-10-18 10:59:51 +09:00

642 lines
20 KiB
Python

import _csv
import csv
import io
from testutils import assert_raises
for row in csv.reader(["one,two,three"]):
[one, two, three] = row
assert one == "one"
assert two == "two"
assert three == "three"
def f():
iter = ["one,two,three", "four,five,six"]
reader = csv.reader(iter)
[one, two, three] = next(reader)
[four, five, six] = next(reader)
assert one == "one"
assert two == "two"
assert three == "three"
assert four == "four"
assert five == "five"
assert six == "six"
f()
def test_delim():
iter = ["one|two|three", "four|five|six"]
reader = csv.reader(iter, delimiter="|")
[one, two, three] = next(reader)
[four, five, six] = next(reader)
assert one == "one"
assert two == "two"
assert three == "three"
assert four == "four"
assert five == "five"
assert six == "six"
with assert_raises(TypeError):
iter = ["one,,two,,three"]
csv.reader(iter, delimiter=",,")
test_delim()
def test_quote_strings_and_notnull_writer():
string_buf = io.StringIO()
csv.writer(string_buf, quoting=csv.QUOTE_STRINGS).writerow(["x", 1, None, ""])
assert string_buf.getvalue() == '"x",1,,""\r\n'
notnull_buf = io.StringIO()
csv.writer(notnull_buf, quoting=csv.QUOTE_NOTNULL).writerow(["x", 1, None, ""])
assert notnull_buf.getvalue() == '"x","1",,""\r\n'
for quoting in (csv.QUOTE_STRINGS, csv.QUOTE_NOTNULL):
buf = io.StringIO()
csv.writer(buf, quoting=quoting).writerow([None, None])
assert buf.getvalue() == ",\r\n"
with assert_raises(csv.Error):
csv.writer(io.StringIO(), quoting=quoting).writerow([None])
with assert_raises(TypeError):
csv.writer(io.StringIO(), quoting=quoting, quotechar=None)
test_quote_strings_and_notnull_writer()
def test_quote_none_writer_without_quotechar():
no_quotechar_buf = io.StringIO()
csv.writer(
no_quotechar_buf,
quoting=csv.QUOTE_NONE,
quotechar=None,
escapechar="\\",
).writerow(["a,b", 'x"y'])
assert no_quotechar_buf.getvalue() == 'a\\,b,x"y\r\n'
default_quotechar_buf = io.StringIO()
csv.writer(
default_quotechar_buf,
quoting=csv.QUOTE_NONE,
escapechar="\\",
).writerow(["a,b", 'x"y'])
assert default_quotechar_buf.getvalue() == 'a\\,b,x\\"y\r\n'
escapechar_buf = io.StringIO()
csv.writer(
escapechar_buf,
quoting=csv.QUOTE_NONE,
quotechar=None,
escapechar="\\",
).writerow(["a\\b"])
assert escapechar_buf.getvalue() == "a\\\\b\r\n"
linebreak_buf = io.StringIO()
csv.writer(
linebreak_buf,
quoting=csv.QUOTE_NONE,
quotechar=None,
escapechar="\\",
).writerow(["a\rb", "c\nd"])
assert linebreak_buf.getvalue() == "a\\\rb,c\\\nd\r\n"
with assert_raises(csv.Error):
csv.writer(io.StringIO(), quoting=csv.QUOTE_NONE, quotechar=None).writerow(
["a,b"]
)
with assert_raises(csv.Error):
csv.writer(
io.StringIO(),
quoting=csv.QUOTE_NONE,
quotechar=None,
escapechar="\\",
).writerow([None])
two_empty_buf = io.StringIO()
csv.writer(
two_empty_buf,
quoting=csv.QUOTE_NONE,
quotechar=None,
escapechar="\\",
).writerow([None, ""])
assert two_empty_buf.getvalue() == ",\r\n"
test_quote_none_writer_without_quotechar()
def test_quote_none_reader_skipinitialspace_escapechar():
reader = csv.reader(
["a, b,\\ c,d"],
quoting=csv.QUOTE_NONE,
escapechar="\\",
skipinitialspace=True,
)
assert list(reader) == [["a", "b", " c", "d"]]
test_quote_none_reader_skipinitialspace_escapechar()
def test_quote_minimal_writer_lineterminator():
# https://github.com/RustPython/RustPython/issues/8302
# QUOTE_MINIMAL must quote '\r' and '\n' regardless of the line terminator.
buf = io.StringIO()
writer = csv.writer(buf, lineterminator="!")
writer.writerow(["a", "b"])
writer.writerow([1, 2])
writer.writerow(["\r", "\n"])
assert buf.getvalue() == 'a,b!1,2!"\r","\n"!'
nul = io.StringIO()
csv.writer(nul, lineterminator="\0").writerow(["\r", "\n"])
assert nul.getvalue() == '"\r","\n"\0'
crlf = io.StringIO()
csv.writer(crlf, lineterminator="!").writerow(["\r\n"])
assert crlf.getvalue() == '"\r\n"!'
# the terminator character itself still triggers quoting
term = io.StringIO()
csv.writer(term, lineterminator="!").writerow(["a!b", "c"])
assert term.getvalue() == '"a!b",c!'
# default terminator behavior is unchanged
default = io.StringIO()
csv.writer(default).writerow(["\r", "\n"])
assert default.getvalue() == '"\r","\n"\r\n'
test_quote_minimal_writer_lineterminator()
def test_multichar_lineterminator():
# https://github.com/RustPython/RustPython/issues/8322
# The writer must store and emit a full multi-character line terminator.
for lineterminator in "\r\n", "\n", "\r", "!@#", "\0":
buf = io.StringIO()
writer = csv.writer(buf, lineterminator=lineterminator)
writer.writerow(["a", "b"])
writer.writerow([1, 2])
writer.writerow(["\r", "\n"])
assert buf.getvalue() == (
f'a,b{lineterminator}1,2{lineterminator}"\r","\n"{lineterminator}'
), (lineterminator, buf.getvalue())
# A field is quoted when it contains any byte of the terminator (QUOTE_MINIMAL).
for field, expected in [
("a@b", '"a@b",x!@#'),
("a!b", '"a!b",x!@#'),
("a#b", '"a#b",x!@#'),
("abc", "abc,x!@#"),
]:
buf = io.StringIO()
csv.writer(buf, lineterminator="!@#").writerow([field, "x"])
assert buf.getvalue() == expected, (field, buf.getvalue())
# The csv-core-backed QUOTE_ALL / QUOTE_NONNUMERIC paths emit the full
# terminator too, and keep the state machine correct across rows.
allq = io.StringIO()
writer = csv.writer(allq, lineterminator="!@#", quoting=csv.QUOTE_ALL)
writer.writerow(["a", "b"])
writer.writerow(["c", "d"])
assert allq.getvalue() == '"a","b"!@#"c","d"!@#', allq.getvalue()
nonnum = io.StringIO()
csv.writer(nonnum, lineterminator="!@#", quoting=csv.QUOTE_NONNUMERIC).writerow(
["a", 1]
)
assert nonnum.getvalue() == '"a",1!@#', nonnum.getvalue()
# A field that itself contains a line-break byte must be kept intact: the
# csv-core path drops only the trailing record terminator, not a byte from
# the field data.
embedded = io.StringIO()
csv.writer(embedded, lineterminator="!@#", quoting=csv.QUOTE_ALL).writerow(
["x\ny", "z"]
)
assert embedded.getvalue() == '"x\ny","z"!@#', embedded.getvalue()
# QUOTE_NONE escapes any byte of the terminator.
none = io.StringIO()
csv.writer(
none, lineterminator="!@#", quoting=csv.QUOTE_NONE, escapechar="\\"
).writerow(["a!b", "x"])
assert none.getvalue() == "a\\!b,x!@#", none.getvalue()
# register_dialect round-trips a multi-character terminator.
csv.register_dialect("multichar_lt", delimiter=",", lineterminator="!@#")
try:
reg = io.StringIO()
csv.writer(reg, dialect="multichar_lt").writerow(["a", "b"])
assert reg.getvalue() == "a,b!@#", reg.getvalue()
finally:
csv.unregister_dialect("multichar_lt")
# The dialect attribute reflects the full terminator.
assert (
csv.writer(io.StringIO(), lineterminator="!@#").dialect.lineterminator == "!@#"
)
# The reader ignores lineterminator (like CPython) and only splits on \r\n.
assert list(csv.reader(io.StringIO("a,b!@#c,d!@#"), lineterminator="!@#")) == [
["a", "b!@#c", "d!@#"]
]
assert list(csv.reader(io.StringIO("a,b\r\nc,d\r\n"), lineterminator="!@#")) == [
["a", "b"],
["c", "d"],
]
test_multichar_lineterminator()
def test_empty_lineterminator():
class EmptyLineTerminator(csv.excel):
lineterminator = ""
keyword = io.StringIO()
csv.writer(keyword, lineterminator="").writerows([["a", "b"], ["c", "d"]])
assert keyword.getvalue() == "a,bc,d"
dialect = io.StringIO()
csv.writer(dialect, dialect=EmptyLineTerminator).writerows([["a", "b"], ["c", "d"]])
assert dialect.getvalue() == "a,bc,d"
source = "a,b\r\nc,d\n"
assert list(csv.reader(io.StringIO(source), lineterminator="")) == [
["a", "b"],
["c", "d"],
]
assert list(csv.reader(io.StringIO(source), dialect=EmptyLineTerminator)) == [
["a", "b"],
["c", "d"],
]
test_empty_lineterminator()
def test_quote_minimal_writer_empty_fields():
buf = io.StringIO()
writer = csv.writer(buf)
writer.writerow([""])
writer.writerow([None])
writer.writerow([])
writer.writerow(["", ""])
assert buf.getvalue() == '""\r\n""\r\n\r\n,\r\n'
test_quote_minimal_writer_empty_fields()
def test_reader_skipinitialspace_preserves_quoted_spaces():
reader = csv.reader(['a, "b, c", d'], skipinitialspace=True)
assert list(reader) == [["a", "b, c", "d"]]
test_reader_skipinitialspace_preserves_quoted_spaces()
# The reader is built with the default dialect, which it has to look up itself.
with assert_raises(StopIteration):
next(_csv.reader([]))
def test_quote_nonnumeric_writer():
class CustomInt:
def __int__(self):
return 42
def __str__(self):
return "42"
class NumberWithDelimiter:
def __int__(self):
return 1
def __str__(self):
return "1,2"
buf = io.StringIO()
writer = csv.writer(buf, quoting=csv.QUOTE_NONNUMERIC)
writer.writerow(["123", 123, "a"])
writer.writerow([12.5, "12.5", True, False, None])
writer.writerow([CustomInt(), NumberWithDelimiter()])
writer.writerow([""])
writer.writerow([None])
writer.writerow([])
writer.writerow(["", ""])
assert buf.getvalue() == (
'"123",123,"a"\r\n'
'12.5,"12.5",True,False,""\r\n'
'42,"1,2"\r\n'
'""\r\n'
'""\r\n'
"\r\n"
'"",""\r\n'
), repr(buf.getvalue())
# QUOTE_NONNUMERIC with space delimiter and skipinitialspace
sp_buf = io.StringIO()
csv.writer(
sp_buf,
delimiter=" ",
skipinitialspace=True,
quoting=csv.QUOTE_NONNUMERIC,
).writerow(["a", "", "b"])
assert sp_buf.getvalue() == '"a" "" "b"\r\n', repr(sp_buf.getvalue())
test_quote_nonnumeric_writer()
def test_dialect_truth_conversion():
def registered(dialect, **kwargs):
name = "truth_conversion"
try:
csv.register_dialect(name, dialect, **kwargs)
return csv.get_dialect(name)
finally:
if name in csv.list_dialects():
csv.unregister_dialect(name)
constructors = (
lambda dialect, **kwargs: csv.reader([], dialect, **kwargs).dialect,
lambda dialect, **kwargs: csv.writer(io.StringIO(), dialect, **kwargs).dialect,
_csv.Dialect,
registered,
)
class TruthError:
def __init__(self, error):
self.error = error
self.calls = 0
def __bool__(self):
self.calls += 1
raise self.error
class InvalidTruth:
def __bool__(self):
return 1
for construct in constructors:
for field in ("strict", "doublequote", "skipinitialspace"):
for error_type in (RuntimeError, AttributeError, TypeError):
for use_class in (False, True):
error = error_type("dialect truth conversion")
value = TruthError(error)
class Dialect(csv.excel):
def __new__(cls):
raise AssertionError(
"dialect classes must not be instantiated"
)
dialect = Dialect if use_class else csv.excel()
setattr(dialect, field, value)
try:
construct(dialect)
except error_type as caught:
assert caught is error
else:
raise AssertionError("dialect truth error was suppressed")
assert value.calls == 1
# An explicit option replaces the attribute before conversion.
result = construct(dialect, **{field: True})
assert getattr(result, field) is True
assert value.calls == 1
value = TruthError(error)
try:
construct(None, **{field: value})
except error_type as caught:
assert caught is error
else:
raise AssertionError("keyword truth error was replaced")
assert value.calls == 1
for value in (None, False, True, 0, 1, [], [1]):
dialect = csv.excel()
setattr(dialect, field, value)
assert getattr(construct(dialect), field) is bool(value)
assert getattr(construct(None, **{field: value}), field) is bool(value)
dialect = csv.excel()
setattr(dialect, field, InvalidTruth())
with assert_raises(TypeError):
construct(dialect)
with assert_raises(TypeError):
construct(None, **{field: InvalidTruth()})
# csv.Dialect's Python validation wrapper translates TypeError to csv.Error.
for error_type in (RuntimeError, AttributeError, TypeError):
error = error_type("direct Python dialect validation")
value = TruthError(error)
class Dialect(csv.excel):
strict = value
expected = csv.Error if error_type is TypeError else error_type
try:
Dialect()
except expected as caught:
assert str(caught) == str(error)
if error_type is not TypeError:
assert caught is error
else:
raise AssertionError("Python dialect validation swallowed the error")
assert value.calls == 1
test_dialect_truth_conversion()
def test_dialect_lookup_and_conversion_order():
fields = (
"delimiter",
"doublequote",
"escapechar",
"lineterminator",
"quotechar",
"quoting",
"skipinitialspace",
"strict",
)
defaults = (",", True, None, "\r\n", '"', csv.QUOTE_MINIMAL, False, False)
constructors = (
lambda dialect, **kwargs: csv.reader([], dialect, **kwargs).dialect,
lambda dialect, **kwargs: csv.writer(io.StringIO(), dialect, **kwargs).dialect,
_csv.Dialect,
)
for construct in constructors:
for error_type in (AttributeError, RuntimeError, KeyboardInterrupt):
events = []
class Missing:
def __getattribute__(self, name):
events.append(name)
raise error_type(name)
if error_type is AttributeError:
result = construct(Missing())
assert tuple(getattr(result, field) for field in fields) == defaults
assert events == list(fields)
else:
with assert_raises(error_type) as caught:
construct(Missing())
assert caught.exception.args == ("delimiter",)
assert events == ["delimiter"]
for dialect in (None, object(), _csv.Dialect(csv.excel)):
result = construct(dialect)
assert tuple(getattr(result, field) for field in fields) == defaults
events = []
class Truth:
def __init__(self, name):
self.name = name
def __bool__(self):
events.append("bool:" + self.name)
return False
values = dict(zip(fields, defaults))
for field in ("doublequote", "skipinitialspace", "strict"):
values[field] = Truth(field)
class Dialect:
def __getattribute__(self, name):
events.append(name)
return values[name]
construct(Dialect())
assert events == list(fields) + [
"bool:doublequote",
"bool:skipinitialspace",
"bool:strict",
]
events.clear()
construct(Dialect(), strict=True, doublequote=True)
assert events == [
field for field in fields if field not in ("strict", "doublequote")
] + ["bool:skipinitialspace"]
events.clear()
construct(
None,
strict=Truth("strict"),
skipinitialspace=Truth("skipinitialspace"),
doublequote=Truth("doublequote"),
)
assert events == ["bool:doublequote", "bool:skipinitialspace", "bool:strict"]
# Collect every attribute before starting conversions, even on failure.
events.clear()
values["delimiter"] = "too long"
with assert_raises(TypeError):
construct(Dialect())
assert events == list(fields)
error = RuntimeError("truth conversion wins")
class BadTruth:
def __bool__(self):
raise error
# Conversion order differs from option-validation order.
for options in (
{"doublequote": BadTruth(), "escapechar": "too long"},
{"strict": BadTruth(), "quoting": 42},
{"strict": BadTruth(), "quotechar": None, "quoting": csv.QUOTE_ALL},
{"strict": BadTruth(), "delimiter": "\n"},
):
try:
construct(None, **options)
except RuntimeError as caught:
assert caught is error
else:
raise AssertionError("conversion order did not preserve truth error")
for options in (
{"delimiter": "too long", "doublequote": BadTruth()},
{"lineterminator": None, "skipinitialspace": BadTruth()},
{"quoting": None, "strict": BadTruth()},
{"unknown_option": None, "strict": BadTruth()},
):
with assert_raises(TypeError):
construct(None, **options)
class IntegerSubclass(int):
pass
class Index:
def __index__(self):
raise AssertionError("quoting must not call __index__")
for value in (True, IntegerSubclass(0), Index()):
with assert_raises(TypeError):
construct(None, quoting=value)
with assert_raises(OverflowError):
construct(None, quoting=2**40, strict=BadTruth())
# None quotechar disables quoting only when quoting was not supplied.
assert construct(None, quotechar=None).quoting == csv.QUOTE_NONE
with assert_raises(TypeError):
construct(csv.excel, quotechar=None)
with assert_raises(TypeError):
construct(None, quotechar=None, quoting=csv.QUOTE_MINIMAL)
assert _csv.Dialect().strict is False
assert _csv.Dialect(strict=True).strict is True
for constructor, arg in ((csv.reader, []), (csv.writer, io.StringIO())):
with assert_raises(TypeError):
constructor(arg, None, None, strict=BadTruth())
test_dialect_lookup_and_conversion_order()
def test_field_size_limit_types():
index_calls = []
class IntSubclass(int):
pass
class Index:
def __index__(self):
index_calls.append(True)
return 17
old_limit = csv.field_size_limit()
try:
assert csv.field_size_limit(23) == old_limit
for value in (True, IntSubclass(17), Index()):
with assert_raises(TypeError):
csv.field_size_limit(value)
assert csv.field_size_limit() == 23
assert index_calls == []
with assert_raises(OverflowError):
csv.field_size_limit(2**100)
assert csv.field_size_limit() == 23
finally:
csv.field_size_limit(old_limit)
test_field_size_limit_types()