FFFF Issue #23171: csv.Writer.writerow() now supports arbitrary iterables. · python/cpython@7901b48 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7901b48

Browse files
serhiy-storchakaserhiy-storchaka
committed
Issue #23171: csv.Writer.writerow() now supports arbitrary iterables.
1 parent a695f83 commit 7901b48

File tree

5 files changed

+54
-46
lines changed

5 files changed

+54
-46
lines changed

‎Doc/library/csv.rst‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -419,7 +419,7 @@ Writer Objects
419419

420420
:class:`Writer` objects (:class:`DictWriter` instances and objects returned by
421421
the :func:`writer` function) have the following public methods. A *row* must be
422-
a sequence of strings or numbers for :class:`Writer` objects and a dictionary
422+
an iterable of strings or numbers for :class:`Writer` objects and a dictionary
423423
mapping fieldnames to strings or numbers (by passing them through :func:`str`
424424
first) for :class:`DictWriter` objects. Note that complex numbers are written
425425
out surrounded by parens. This may cause some problems for other programs which
@@ -431,6 +431,8 @@ read CSV files (assuming they support complex numbers at all).
431431
Write the *row* parameter to the writer's file object, formatted according to
432432
the current dialect.
433433

434+
.. versionchanged:: 3.5
435+
Added support of arbitrary iterables.
434436

435437
.. method:: csvwriter.writerows(rows)
436438

‎Lib/csv.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -147,16 +147,13 @@ def _dict_to_list(self, rowdict):
147147
if wrong_fields:
148148
raise ValueError("dict contains fields not in fieldnames: "
149149
+ ", ".join([repr(x) for x in wrong_fields]))
150-
return [rowdict.get(key, self.restval) for key in self.fieldnames]
150+
return (rowdict.get(key, self.restval) for key in self.fieldnames)
151151

152152
def writerow(self, rowdict):
153153
return self.writer.writerow(self._dict_to_list(rowdict))
154154

155155
def writerows(self, rowdicts):
156-
rows = []
157-
for rowdict in rowdicts:
158-
rows.append(self._dict_to_list(rowdict))
159-
return self.writer.writerows(rows)
156+
return self.writer.writerows(map(self._dict_to_list, rowdicts))
160157

161158
# Guard Sniffer's type checking against builds that exclude complex()
162159
try:

‎Lib/test/test_csv.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,14 @@ def test_write_escape(self):
186186
self._write_test(['a',1,'p,q'], 'a,1,p\\,q',
187187
escapechar='\\', quoting = csv.QUOTE_NONE)
188188

189+
def test_write_iterable(self):
190+
self._write_test(iter(['a', 1, 'p,q']), 'a,1,"p,q"')
191+
self._write_test(iter(['a', 1, None]), 'a,1,')
192+
self._write_test(iter([]), '')
193+
self._write_test(iter([None]), '""')
194+
self._write_error_test(csv.Error, iter([None]), quoting=csv.QUOTE_NONE)
195+
self._write_test(iter([None, None]), ',')
196+
189197
def test_writerows(self):
190198
class BrokenFile:
191199
def write(self, buf):

‎Misc/NEWS‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ Core and Builtins
5656
Library
5757
-------
5858

59+
- Issue #23171: csv.Writer.writerow() now supports arbitrary iterables.
60+
5961
- Issue #23745: The new email header parser now handles duplicate MIME
6062
parameter names without error, similar to how get_param behaves.
6163

‎Modules/_csv.c‎

Lines changed: 39 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1009,7 +1009,7 @@ join_reset(WriterObj *self)
10091009
*/
10101010
static Py_ssize_t
10111011
join_append_data(WriterObj *self, unsigned int field_kind, void *field_data,
1012-
Py_ssize_t field_len, int quote_empty, int *quoted,
1012+
Py_ssize_t field_len, int *quoted,
10131013
int copy_phase)
10141014
{
10151015
DialectObj *dialect = self->dialect;
@@ -1071,18 +1071,6 @@ join_append_data(WriterObj *self, unsigned int field_kind, void *field_data,
10711071
ADDCH(c);
10721072
}
10731073

1074-
/* If field is empty check if it needs to be quoted.
1075-
*/
1076-
if (i == 0 && quote_empty) {
1077-
if (dialect->quoting == QUOTE_NONE) {
1078-
PyErr_Format(_csvstate_global->error_obj,
1079-
"single empty field record must be quoted");
1080-
return -1;
1081-
}
1082-
else
1083-
*quoted = 1;
1084-
}
1085-
10861074
if (*quoted) {
10871075
if (copy_phase)
10881076
ADDCH(dialect->quotechar);
@@ -1126,7 +1114,7 @@ join_check_rec_size(WriterObj *self, Py_ssize_t rec_len)
11261114
}
11271115

11281116
static int
1129-
join_append(WriterObj *self, PyObject *field, int *quoted, int quote_empty)
1117+
join_append(WriterObj *self, PyObject *field, int quoted)
11301118
{
11311119
unsigned int field_kind = -1;
11321120
void *field_data = NULL;
@@ -1141,7 +1129,7 @@ join_append(WriterObj *self, PyObject *field, int *quoted, int quote_empty)
11411129
field_len = PyUnicode_GET_LENGTH(field);
11421130
}
11431131
rec_len = join_append_data(self, field_kind, field_data, field_len,
1144-
quote_empty, quoted, 0);
1132+
&quoted, 0);
11451133
if (rec_len < 0)
11461134
return 0;
11471135

@@ -1150,7 +1138,7 @@ join_append(WriterObj *self, PyObject *field, int *quoted, int quote_empty)
11501138
return 0;
11511139

11521140
self->rec_len = join_append_data(self, field_kind, field_data, field_len,
1153-
quote_empty, quoted, 1);
1141+
&quoted, 1);
11541142
self->num_fields++;
11551143

11561144
return 1;
@@ -1181,37 +1169,30 @@ join_append_lineterminator(WriterObj *self)
11811169
}
11821170

11831171
PyDoc_STRVAR(csv_writerow_doc,
1184-
"writerow(sequence)\n"
1172+
"writerow(iterable)\n"
11851173
"\n"
1186-
"Construct and write a CSV record from a sequence of fields. Non-string\n"
1174+
"Construct and write a CSV record from an iterable of fields. Non-string\n"
11871175
"elements will be converted to string.");
11881176

11891177
static PyObject *
11901178
csv_writerow(WriterObj *self, PyObject *seq)
11911179
{
11921180
DialectObj *dialect = self->dialect;
1193-
Py_ssize_t len, i;
1194-
PyObject *line, *result;
1181+
PyObject *iter, *field, *line, *result;
11951182

1196-
if (!PySequence_Check(seq))
1197-
return PyErr_Format(_csvstate_global->error_obj, "sequence expected");
1198-
1199-
len = PySequence_Length(seq);
1200-
if (len < 0)
1201-
return NULL;
1183+
iter = PyObject_GetIter(seq);
1184+
if (iter == NULL)
1185+
return PyErr_Format(_csvstate_global->error_obj,
1186+
"iterable expected, not %.200s",
1187+
seq->ob_type->tp_name);
12021188

12031189
/* Join all fields in internal buffer.
12041190
*/
12051191
join_reset(self);
1206-
for (i = 0; i < len; i++) {
1207-
PyObject *field;
1192+
while ((field = PyIter_Next(iter))) {
12081193
int append_ok;
12091194
int quoted;
12101195

1211-
field = PySequence_GetItem(seq, i);
1212-
if (field == NULL)
1213-
return NULL;
1214-
12151196
switch (dialect->quoting) {
12161197
case QUOTE_NONNUMERIC:
12171198
quoted = !PyNumber_Check(field);
@@ -1225,31 +1206,49 @@ csv_writerow(WriterObj *self, PyObject *seq)
12251206
}
12261207

12271208
if (PyUnicode_Check(field)) {
1228-
append_ok = join_append(self, field, &quoted, len == 1);
1209+
append_ok = join_append(self, field, quoted);
12291210
Py_DECREF(field);
12301211
}
12311212
else if (field == Py_None) {
1232-
append_ok = join_append(self, NULL, &quoted, len == 1);
1213+
append_ok = join_append(self, NULL, quoted);
12331214
Py_DECREF(field);
12341215
}
12351216
else {
12361217
PyObject *str;
12371218

12381219
str = PyObject_Str(field);
12391220
Py_DECREF(field);
1240-
if (str == NULL)
1221+
if (str == NULL) {
1222+
Py_DECREF(iter);
12411223
return NULL;
1242-
append_ok = join_append(self, str, &quoted, len == 1);
1224+
}
1225+
append_ok = join_append(self, str, quoted);
12431226
Py_DECREF(str);
12441227
}
1245-
if (!append_ok)
1228+
if (!append_ok) {
1229+
Py_DECREF(iter);
1230+
return NULL;
1231+
}
1232+
}
1233+
Py_DECREF(iter);
1234+
if (PyErr_Occurred())
1235+
return NULL;
1236+
1237+
if (self->num_fields > 0 && self->rec_size == 0) {
1238+
if (dialect->quoting == QUOTE_NONE) {
1239+
PyErr_Format(_csvstate_global->error_obj,
1240+
"single empty field record must be quoted");
1241+
return NULL;
1242+
}
1243+
self->num_fields--;
1244+
if (!join_append(self, NULL, 1))
12461245
return NULL;
12471246
}
12481247

12491248
/* Add line terminator.
12501249
*/
12511250
if (!join_append_lineterminator(self))
1252-
return 0;
1251+
return NULL;
12531252

12541253
line = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND,
12551254
(void *) self->rec, self->rec_len);
@@ -1261,9 +1260,9 @@ csv_writerow(WriterObj *self, PyObject *seq)
12611260
}
12621261

12631262
PyDoc_STRVAR(csv_writerows_doc,
1264-
"writerows(sequence of sequences)\n"
1263+
"writerows(iterable of iterables)\n"
12651264
"\n"
1266-
"Construct and write a series of sequences to a csv file. Non-string\n"
1265+
"Construct and write a series of iterables to a csv file. Non-string\n"
12671266
"elements will be converted to string.");
12681267

12691268
static PyObject *

0 commit comments

Comments
 (0)
0