blob: 48a5cf809420129566b90d7ccad0e5a642745bf1 [file] [log] [blame]
Skip Montanaroa16b21f2003-03-23 14:32:54 +00001/* csv module */
2
3/*
4
5This module provides the low-level underpinnings of a CSV reading/writing
6module. Users should not use this module directly, but import the csv.py
7module instead.
8
Skip Montanarob4a04172003-03-20 23:29:12 +00009*/
10
Skip Montanaro7b01a832003-04-12 19:23:46 +000011#define MODULE_VERSION "1.0"
12
Skip Montanarob4a04172003-03-20 23:29:12 +000013#include "Python.h"
14#include "structmember.h"
15
Andrew McNamara37d2bdf2005-01-10 12:22:48 +000016
Antoine Pitroue7672d32012-05-16 11:33:08 +020017typedef struct {
18 PyObject *error_obj; /* CSV exception */
19 PyObject *dialects; /* Dialect registry */
20 long field_limit; /* max parsed field size */
21} _csvstate;
22
23#define _csvstate(o) ((_csvstate *)PyModule_GetState(o))
24
25static int
26_csv_clear(PyObject *m)
27{
28 Py_CLEAR(_csvstate(m)->error_obj);
29 Py_CLEAR(_csvstate(m)->dialects);
30 return 0;
31}
32
33static int
34_csv_traverse(PyObject *m, visitproc visit, void *arg)
35{
36 Py_VISIT(_csvstate(m)->error_obj);
37 Py_VISIT(_csvstate(m)->dialects);
38 return 0;
39}
40
41static void
42_csv_free(void *m)
43{
44 _csv_clear((PyObject *)m);
45}
46
47static struct PyModuleDef _csvmodule;
48
49#define _csvstate_global ((_csvstate *)PyModule_GetState(PyState_FindModule(&_csvmodule)))
Skip Montanarob4a04172003-03-20 23:29:12 +000050
51typedef enum {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000052 START_RECORD, START_FIELD, ESCAPED_CHAR, IN_FIELD,
53 IN_QUOTED_FIELD, ESCAPE_IN_QUOTED_FIELD, QUOTE_IN_QUOTED_FIELD,
54 EAT_CRNL
Skip Montanarob4a04172003-03-20 23:29:12 +000055} ParserState;
56
57typedef enum {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000058 QUOTE_MINIMAL, QUOTE_ALL, QUOTE_NONNUMERIC, QUOTE_NONE
Skip Montanarob4a04172003-03-20 23:29:12 +000059} QuoteStyle;
60
61typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000062 QuoteStyle style;
63 char *name;
Skip Montanarob4a04172003-03-20 23:29:12 +000064} StyleDesc;
65
66static StyleDesc quote_styles[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000067 { QUOTE_MINIMAL, "QUOTE_MINIMAL" },
68 { QUOTE_ALL, "QUOTE_ALL" },
69 { QUOTE_NONNUMERIC, "QUOTE_NONNUMERIC" },
70 { QUOTE_NONE, "QUOTE_NONE" },
71 { 0 }
Skip Montanarob4a04172003-03-20 23:29:12 +000072};
73
74typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000075 PyObject_HEAD
Guido van Rossum46264582007-08-06 19:32:18 +000076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000077 int doublequote; /* is " represented by ""? */
Antoine Pitrou77ea6402011-10-07 04:26:55 +020078 Py_UCS4 delimiter; /* field separator */
79 Py_UCS4 quotechar; /* quote character */
80 Py_UCS4 escapechar; /* escape character */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 int skipinitialspace; /* ignore spaces following delimiter? */
82 PyObject *lineterminator; /* string to write between records */
83 int quoting; /* style of quoting to write */
Skip Montanarob4a04172003-03-20 23:29:12 +000084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000085 int strict; /* raise exception on bad CSV */
Skip Montanarob4a04172003-03-20 23:29:12 +000086} DialectObj;
87
Neal Norwitz227b5332006-03-22 09:28:35 +000088static PyTypeObject Dialect_Type;
Skip Montanarob4a04172003-03-20 23:29:12 +000089
90typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000091 PyObject_HEAD
Skip Montanarob4a04172003-03-20 23:29:12 +000092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000093 PyObject *input_iter; /* iterate over this for input lines */
Skip Montanarob4a04172003-03-20 23:29:12 +000094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000095 DialectObj *dialect; /* parsing dialect */
Skip Montanarob4a04172003-03-20 23:29:12 +000096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000097 PyObject *fields; /* field list for current record */
98 ParserState state; /* current CSV parse state */
Antoine Pitrou77ea6402011-10-07 04:26:55 +020099 Py_UCS4 *field; /* temporary buffer */
Antoine Pitrou40455752010-08-15 18:51:10 +0000100 Py_ssize_t field_size; /* size of allocated buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000101 Py_ssize_t field_len; /* length of current field */
102 int numeric_field; /* treat field as numeric */
103 unsigned long line_num; /* Source-file line number */
Skip Montanarob4a04172003-03-20 23:29:12 +0000104} ReaderObj;
105
Neal Norwitz227b5332006-03-22 09:28:35 +0000106static PyTypeObject Reader_Type;
Skip Montanarob4a04172003-03-20 23:29:12 +0000107
Christian Heimes90aa7642007-12-19 02:45:37 +0000108#define ReaderObject_Check(v) (Py_TYPE(v) == &Reader_Type)
Skip Montanarob4a04172003-03-20 23:29:12 +0000109
110typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000111 PyObject_HEAD
Skip Montanarob4a04172003-03-20 23:29:12 +0000112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000113 PyObject *writeline; /* write output lines to this file */
Skip Montanarob4a04172003-03-20 23:29:12 +0000114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000115 DialectObj *dialect; /* parsing dialect */
Skip Montanarob4a04172003-03-20 23:29:12 +0000116
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200117 Py_UCS4 *rec; /* buffer for parser.join */
Antoine Pitrou40455752010-08-15 18:51:10 +0000118 Py_ssize_t rec_size; /* size of allocated record */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000119 Py_ssize_t rec_len; /* length of record */
120 int num_fields; /* number of fields in record */
Guido van Rossum46264582007-08-06 19:32:18 +0000121} WriterObj;
Skip Montanarob4a04172003-03-20 23:29:12 +0000122
Neal Norwitz227b5332006-03-22 09:28:35 +0000123static PyTypeObject Writer_Type;
Skip Montanarob4a04172003-03-20 23:29:12 +0000124
125/*
126 * DIALECT class
127 */
128
129static PyObject *
130get_dialect_from_registry(PyObject * name_obj)
131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 PyObject *dialect_obj;
Skip Montanarob4a04172003-03-20 23:29:12 +0000133
Antoine Pitroue7672d32012-05-16 11:33:08 +0200134 dialect_obj = PyDict_GetItem(_csvstate_global->dialects, name_obj);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000135 if (dialect_obj == NULL) {
136 if (!PyErr_Occurred())
Antoine Pitroue7672d32012-05-16 11:33:08 +0200137 PyErr_Format(_csvstate_global->error_obj, "unknown dialect");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000138 }
139 else
140 Py_INCREF(dialect_obj);
141 return dialect_obj;
Skip Montanarob4a04172003-03-20 23:29:12 +0000142}
143
Skip Montanarob4a04172003-03-20 23:29:12 +0000144static PyObject *
145get_string(PyObject *str)
146{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000147 Py_XINCREF(str);
148 return str;
Skip Montanarob4a04172003-03-20 23:29:12 +0000149}
150
Skip Montanarob4a04172003-03-20 23:29:12 +0000151static PyObject *
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200152get_nullchar_as_None(Py_UCS4 c)
Skip Montanarob4a04172003-03-20 23:29:12 +0000153{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000154 if (c == '\0') {
155 Py_INCREF(Py_None);
156 return Py_None;
157 }
158 else
Martin v. Löwisd63a3b82011-09-28 07:41:54 +0200159 return PyUnicode_FromOrdinal(c);
Skip Montanarob4a04172003-03-20 23:29:12 +0000160}
161
Skip Montanarob4a04172003-03-20 23:29:12 +0000162static PyObject *
163Dialect_get_lineterminator(DialectObj *self)
164{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000165 return get_string(self->lineterminator);
Skip Montanarob4a04172003-03-20 23:29:12 +0000166}
167
Skip Montanarob4a04172003-03-20 23:29:12 +0000168static PyObject *
Guido van Rossuma9769c22007-08-07 23:59:30 +0000169Dialect_get_delimiter(DialectObj *self)
170{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000171 return get_nullchar_as_None(self->delimiter);
Guido van Rossuma9769c22007-08-07 23:59:30 +0000172}
173
174static PyObject *
Skip Montanarob4a04172003-03-20 23:29:12 +0000175Dialect_get_escapechar(DialectObj *self)
176{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000177 return get_nullchar_as_None(self->escapechar);
Skip Montanarob4a04172003-03-20 23:29:12 +0000178}
179
Andrew McNamara1196cf12005-01-07 04:42:45 +0000180static PyObject *
181Dialect_get_quotechar(DialectObj *self)
Skip Montanarob4a04172003-03-20 23:29:12 +0000182{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000183 return get_nullchar_as_None(self->quotechar);
Skip Montanarob4a04172003-03-20 23:29:12 +0000184}
185
186static PyObject *
187Dialect_get_quoting(DialectObj *self)
188{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 return PyLong_FromLong(self->quoting);
Skip Montanarob4a04172003-03-20 23:29:12 +0000190}
191
192static int
Andrew McNamara1196cf12005-01-07 04:42:45 +0000193_set_bool(const char *name, int *target, PyObject *src, int dflt)
Skip Montanarob4a04172003-03-20 23:29:12 +0000194{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000195 if (src == NULL)
196 *target = dflt;
Antoine Pitrou6f430e42012-08-15 23:18:25 +0200197 else {
198 int b = PyObject_IsTrue(src);
199 if (b < 0)
200 return -1;
201 *target = b;
202 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000203 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000204}
205
Andrew McNamara1196cf12005-01-07 04:42:45 +0000206static int
207_set_int(const char *name, int *target, PyObject *src, int dflt)
208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000209 if (src == NULL)
210 *target = dflt;
211 else {
212 long value;
213 if (!PyLong_CheckExact(src)) {
214 PyErr_Format(PyExc_TypeError,
215 "\"%s\" must be an integer", name);
216 return -1;
217 }
218 value = PyLong_AsLong(src);
219 if (value == -1 && PyErr_Occurred())
220 return -1;
Martin v. Löwisd1a1d1e2007-12-04 22:10:37 +0000221#if SIZEOF_LONG > SIZEOF_INT
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000222 if (value > INT_MAX || value < INT_MIN) {
223 PyErr_Format(PyExc_ValueError,
224 "integer out of range for \"%s\"", name);
225 return -1;
226 }
Martin v. Löwisd1a1d1e2007-12-04 22:10:37 +0000227#endif
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000228 *target = (int)value;
229 }
230 return 0;
Andrew McNamara1196cf12005-01-07 04:42:45 +0000231}
232
233static int
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200234_set_char(const char *name, Py_UCS4 *target, PyObject *src, Py_UCS4 dflt)
Andrew McNamara1196cf12005-01-07 04:42:45 +0000235{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 if (src == NULL)
237 *target = dflt;
238 else {
239 *target = '\0';
240 if (src != Py_None) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000241 Py_ssize_t len;
Victor Stinner9e30aa52011-11-21 02:49:52 +0100242 len = PyUnicode_GetLength(src);
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200243 if (len > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000244 PyErr_Format(PyExc_TypeError,
245 "\"%s\" must be an 1-character string",
246 name);
247 return -1;
248 }
Stefan Krahe6996ed2012-11-02 14:44:20 +0100249 /* PyUnicode_READY() is called in PyUnicode_GetLength() */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000250 if (len > 0)
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200251 *target = PyUnicode_READ_CHAR(src, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000252 }
253 }
254 return 0;
Andrew McNamara1196cf12005-01-07 04:42:45 +0000255}
256
257static int
258_set_str(const char *name, PyObject **target, PyObject *src, const char *dflt)
259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000260 if (src == NULL)
261 *target = PyUnicode_DecodeASCII(dflt, strlen(dflt), NULL);
262 else {
263 if (src == Py_None)
264 *target = NULL;
Stefan Krahe6996ed2012-11-02 14:44:20 +0100265 else if (!PyUnicode_Check(src)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000266 PyErr_Format(PyExc_TypeError,
267 "\"%s\" must be a string", name);
268 return -1;
269 }
270 else {
Stefan Krahe6996ed2012-11-02 14:44:20 +0100271 if (PyUnicode_READY(src) == -1)
272 return -1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000273 Py_XDECREF(*target);
274 Py_INCREF(src);
275 *target = src;
276 }
277 }
278 return 0;
Andrew McNamara1196cf12005-01-07 04:42:45 +0000279}
280
281static int
282dialect_check_quoting(int quoting)
283{
Victor Stinner4fe519b2010-11-09 09:40:16 +0000284 StyleDesc *qs;
Andrew McNamara1196cf12005-01-07 04:42:45 +0000285
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000286 for (qs = quote_styles; qs->name; qs++) {
287 if (qs->style == quoting)
288 return 0;
289 }
290 PyErr_Format(PyExc_TypeError, "bad \"quoting\" value");
291 return -1;
Andrew McNamara1196cf12005-01-07 04:42:45 +0000292}
Skip Montanarob4a04172003-03-20 23:29:12 +0000293
294#define D_OFF(x) offsetof(DialectObj, x)
295
296static struct PyMemberDef Dialect_memberlist[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000297 { "skipinitialspace", T_INT, D_OFF(skipinitialspace), READONLY },
298 { "doublequote", T_INT, D_OFF(doublequote), READONLY },
299 { "strict", T_INT, D_OFF(strict), READONLY },
300 { NULL }
Skip Montanarob4a04172003-03-20 23:29:12 +0000301};
302
303static PyGetSetDef Dialect_getsetlist[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000304 { "delimiter", (getter)Dialect_get_delimiter},
305 { "escapechar", (getter)Dialect_get_escapechar},
306 { "lineterminator", (getter)Dialect_get_lineterminator},
307 { "quotechar", (getter)Dialect_get_quotechar},
308 { "quoting", (getter)Dialect_get_quoting},
309 {NULL},
Skip Montanarob4a04172003-03-20 23:29:12 +0000310};
311
312static void
313Dialect_dealloc(DialectObj *self)
314{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315 Py_XDECREF(self->lineterminator);
316 Py_TYPE(self)->tp_free((PyObject *)self);
Skip Montanarob4a04172003-03-20 23:29:12 +0000317}
318
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +0000319static char *dialect_kws[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000320 "dialect",
321 "delimiter",
322 "doublequote",
323 "escapechar",
324 "lineterminator",
325 "quotechar",
326 "quoting",
327 "skipinitialspace",
328 "strict",
329 NULL
Andrew McNamara1196cf12005-01-07 04:42:45 +0000330};
331
Andrew McNamara29bf4e42005-01-11 04:49:53 +0000332static PyObject *
333dialect_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
Skip Montanarob4a04172003-03-20 23:29:12 +0000334{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000335 DialectObj *self;
336 PyObject *ret = NULL;
337 PyObject *dialect = NULL;
338 PyObject *delimiter = NULL;
339 PyObject *doublequote = NULL;
340 PyObject *escapechar = NULL;
341 PyObject *lineterminator = NULL;
342 PyObject *quotechar = NULL;
343 PyObject *quoting = NULL;
344 PyObject *skipinitialspace = NULL;
345 PyObject *strict = NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +0000346
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000347 if (!PyArg_ParseTupleAndKeywords(args, kwargs,
348 "|OOOOOOOOO", dialect_kws,
349 &dialect,
350 &delimiter,
351 &doublequote,
352 &escapechar,
353 &lineterminator,
354 &quotechar,
355 &quoting,
356 &skipinitialspace,
357 &strict))
358 return NULL;
Andrew McNamara29bf4e42005-01-11 04:49:53 +0000359
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000360 if (dialect != NULL) {
Stefan Krahe6996ed2012-11-02 14:44:20 +0100361 if (PyUnicode_Check(dialect)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000362 dialect = get_dialect_from_registry(dialect);
363 if (dialect == NULL)
364 return NULL;
365 }
366 else
367 Py_INCREF(dialect);
368 /* Can we reuse this instance? */
369 if (PyObject_TypeCheck(dialect, &Dialect_Type) &&
370 delimiter == 0 &&
371 doublequote == 0 &&
372 escapechar == 0 &&
373 lineterminator == 0 &&
374 quotechar == 0 &&
375 quoting == 0 &&
376 skipinitialspace == 0 &&
377 strict == 0)
378 return dialect;
379 }
Andrew McNamara29bf4e42005-01-11 04:49:53 +0000380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000381 self = (DialectObj *)type->tp_alloc(type, 0);
382 if (self == NULL) {
383 Py_XDECREF(dialect);
384 return NULL;
385 }
386 self->lineterminator = NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +0000387
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000388 Py_XINCREF(delimiter);
389 Py_XINCREF(doublequote);
390 Py_XINCREF(escapechar);
391 Py_XINCREF(lineterminator);
392 Py_XINCREF(quotechar);
393 Py_XINCREF(quoting);
394 Py_XINCREF(skipinitialspace);
395 Py_XINCREF(strict);
396 if (dialect != NULL) {
Andrew McNamara1196cf12005-01-07 04:42:45 +0000397#define DIALECT_GETATTR(v, n) \
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 if (v == NULL) \
399 v = PyObject_GetAttrString(dialect, n)
400 DIALECT_GETATTR(delimiter, "delimiter");
401 DIALECT_GETATTR(doublequote, "doublequote");
402 DIALECT_GETATTR(escapechar, "escapechar");
403 DIALECT_GETATTR(lineterminator, "lineterminator");
404 DIALECT_GETATTR(quotechar, "quotechar");
405 DIALECT_GETATTR(quoting, "quoting");
406 DIALECT_GETATTR(skipinitialspace, "skipinitialspace");
407 DIALECT_GETATTR(strict, "strict");
408 PyErr_Clear();
409 }
Skip Montanarob4a04172003-03-20 23:29:12 +0000410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000411 /* check types and convert to C values */
Andrew McNamara1196cf12005-01-07 04:42:45 +0000412#define DIASET(meth, name, target, src, dflt) \
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 if (meth(name, target, src, dflt)) \
414 goto err
415 DIASET(_set_char, "delimiter", &self->delimiter, delimiter, ',');
416 DIASET(_set_bool, "doublequote", &self->doublequote, doublequote, 1);
417 DIASET(_set_char, "escapechar", &self->escapechar, escapechar, 0);
418 DIASET(_set_str, "lineterminator", &self->lineterminator, lineterminator, "\r\n");
419 DIASET(_set_char, "quotechar", &self->quotechar, quotechar, '"');
420 DIASET(_set_int, "quoting", &self->quoting, quoting, QUOTE_MINIMAL);
421 DIASET(_set_bool, "skipinitialspace", &self->skipinitialspace, skipinitialspace, 0);
422 DIASET(_set_bool, "strict", &self->strict, strict, 0);
Skip Montanarob4a04172003-03-20 23:29:12 +0000423
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000424 /* validate options */
425 if (dialect_check_quoting(self->quoting))
426 goto err;
427 if (self->delimiter == 0) {
428 PyErr_SetString(PyExc_TypeError, "delimiter must be set");
429 goto err;
430 }
431 if (quotechar == Py_None && quoting == NULL)
432 self->quoting = QUOTE_NONE;
433 if (self->quoting != QUOTE_NONE && self->quotechar == 0) {
434 PyErr_SetString(PyExc_TypeError,
435 "quotechar must be set if quoting enabled");
436 goto err;
437 }
438 if (self->lineterminator == 0) {
439 PyErr_SetString(PyExc_TypeError, "lineterminator must be set");
440 goto err;
441 }
Andrew McNamara1196cf12005-01-07 04:42:45 +0000442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000443 ret = (PyObject *)self;
444 Py_INCREF(self);
Andrew McNamara1196cf12005-01-07 04:42:45 +0000445err:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000446 Py_XDECREF(self);
447 Py_XDECREF(dialect);
448 Py_XDECREF(delimiter);
449 Py_XDECREF(doublequote);
450 Py_XDECREF(escapechar);
451 Py_XDECREF(lineterminator);
452 Py_XDECREF(quotechar);
453 Py_XDECREF(quoting);
454 Py_XDECREF(skipinitialspace);
455 Py_XDECREF(strict);
456 return ret;
Skip Montanarob4a04172003-03-20 23:29:12 +0000457}
458
459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000460PyDoc_STRVAR(Dialect_Type_doc,
Skip Montanarob4a04172003-03-20 23:29:12 +0000461"CSV dialect\n"
462"\n"
463"The Dialect type records CSV parsing and generation options.\n");
464
465static PyTypeObject Dialect_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000466 PyVarObject_HEAD_INIT(NULL, 0)
467 "_csv.Dialect", /* tp_name */
468 sizeof(DialectObj), /* tp_basicsize */
469 0, /* tp_itemsize */
470 /* methods */
471 (destructor)Dialect_dealloc, /* tp_dealloc */
472 (printfunc)0, /* tp_print */
473 (getattrfunc)0, /* tp_getattr */
474 (setattrfunc)0, /* tp_setattr */
475 0, /* tp_reserved */
476 (reprfunc)0, /* tp_repr */
477 0, /* tp_as_number */
478 0, /* tp_as_sequence */
479 0, /* tp_as_mapping */
480 (hashfunc)0, /* tp_hash */
481 (ternaryfunc)0, /* tp_call */
482 (reprfunc)0, /* tp_str */
483 0, /* tp_getattro */
484 0, /* tp_setattro */
485 0, /* tp_as_buffer */
486 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
487 Dialect_Type_doc, /* tp_doc */
488 0, /* tp_traverse */
489 0, /* tp_clear */
490 0, /* tp_richcompare */
491 0, /* tp_weaklistoffset */
492 0, /* tp_iter */
493 0, /* tp_iternext */
494 0, /* tp_methods */
495 Dialect_memberlist, /* tp_members */
496 Dialect_getsetlist, /* tp_getset */
497 0, /* tp_base */
498 0, /* tp_dict */
499 0, /* tp_descr_get */
500 0, /* tp_descr_set */
501 0, /* tp_dictoffset */
502 0, /* tp_init */
503 0, /* tp_alloc */
504 dialect_new, /* tp_new */
505 0, /* tp_free */
Skip Montanarob4a04172003-03-20 23:29:12 +0000506};
507
Andrew McNamara91b97462005-01-11 01:07:23 +0000508/*
509 * Return an instance of the dialect type, given a Python instance or kwarg
510 * description of the dialect
511 */
512static PyObject *
513_call_dialect(PyObject *dialect_inst, PyObject *kwargs)
514{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000515 PyObject *ctor_args;
516 PyObject *dialect;
Andrew McNamara91b97462005-01-11 01:07:23 +0000517
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000518 ctor_args = Py_BuildValue(dialect_inst ? "(O)" : "()", dialect_inst);
519 if (ctor_args == NULL)
520 return NULL;
521 dialect = PyObject_Call((PyObject *)&Dialect_Type, ctor_args, kwargs);
522 Py_DECREF(ctor_args);
523 return dialect;
Andrew McNamara91b97462005-01-11 01:07:23 +0000524}
525
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000526/*
527 * READER
528 */
Andrew McNamara0f0599d2005-01-12 09:45:18 +0000529static int
Skip Montanarob4a04172003-03-20 23:29:12 +0000530parse_save_field(ReaderObj *self)
531{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000532 PyObject *field;
Skip Montanarob4a04172003-03-20 23:29:12 +0000533
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200534 field = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND,
535 (void *) self->field, self->field_len);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000536 if (field == NULL)
537 return -1;
538 self->field_len = 0;
539 if (self->numeric_field) {
540 PyObject *tmp;
Andrew McNamara0f0599d2005-01-12 09:45:18 +0000541
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 self->numeric_field = 0;
543 tmp = PyNumber_Float(field);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000544 Py_DECREF(field);
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200545 if (tmp == NULL)
546 return -1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000547 field = tmp;
548 }
549 PyList_Append(self->fields, field);
550 Py_DECREF(field);
551 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000552}
553
554static int
555parse_grow_buff(ReaderObj *self)
556{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000557 if (self->field_size == 0) {
558 self->field_size = 4096;
559 if (self->field != NULL)
560 PyMem_Free(self->field);
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200561 self->field = PyMem_New(Py_UCS4, self->field_size);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000562 }
563 else {
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200564 Py_UCS4 *field = self->field;
Antoine Pitrou40455752010-08-15 18:51:10 +0000565 if (self->field_size > PY_SSIZE_T_MAX / 2) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000566 PyErr_NoMemory();
567 return 0;
568 }
569 self->field_size *= 2;
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200570 self->field = PyMem_Resize(field, Py_UCS4, self->field_size);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000571 }
572 if (self->field == NULL) {
573 PyErr_NoMemory();
574 return 0;
575 }
576 return 1;
Skip Montanarob4a04172003-03-20 23:29:12 +0000577}
578
Andrew McNamarae4d05c42005-01-11 07:32:02 +0000579static int
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200580parse_add_char(ReaderObj *self, Py_UCS4 c)
Skip Montanarob4a04172003-03-20 23:29:12 +0000581{
Antoine Pitroue7672d32012-05-16 11:33:08 +0200582 if (self->field_len >= _csvstate_global->field_limit) {
583 PyErr_Format(_csvstate_global->error_obj, "field larger than field limit (%ld)",
584 _csvstate_global->field_limit);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 return -1;
586 }
587 if (self->field_len == self->field_size && !parse_grow_buff(self))
588 return -1;
589 self->field[self->field_len++] = c;
590 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000591}
592
Andrew McNamarae4d05c42005-01-11 07:32:02 +0000593static int
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200594parse_process_char(ReaderObj *self, Py_UCS4 c)
Skip Montanarob4a04172003-03-20 23:29:12 +0000595{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000596 DialectObj *dialect = self->dialect;
Skip Montanarob4a04172003-03-20 23:29:12 +0000597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000598 switch (self->state) {
599 case START_RECORD:
600 /* start of record */
601 if (c == '\0')
602 /* empty line - return [] */
603 break;
604 else if (c == '\n' || c == '\r') {
605 self->state = EAT_CRNL;
606 break;
607 }
608 /* normal character - handle as START_FIELD */
609 self->state = START_FIELD;
610 /* fallthru */
611 case START_FIELD:
612 /* expecting field */
613 if (c == '\n' || c == '\r' || c == '\0') {
614 /* save empty field - return [fields] */
615 if (parse_save_field(self) < 0)
616 return -1;
617 self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
618 }
619 else if (c == dialect->quotechar &&
620 dialect->quoting != QUOTE_NONE) {
621 /* start quoted field */
622 self->state = IN_QUOTED_FIELD;
623 }
624 else if (c == dialect->escapechar) {
625 /* possible escaped character */
626 self->state = ESCAPED_CHAR;
627 }
628 else if (c == ' ' && dialect->skipinitialspace)
629 /* ignore space at start of field */
630 ;
631 else if (c == dialect->delimiter) {
632 /* save empty field */
633 if (parse_save_field(self) < 0)
634 return -1;
635 }
636 else {
637 /* begin new unquoted field */
638 if (dialect->quoting == QUOTE_NONNUMERIC)
639 self->numeric_field = 1;
640 if (parse_add_char(self, c) < 0)
641 return -1;
642 self->state = IN_FIELD;
643 }
644 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000645
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000646 case ESCAPED_CHAR:
647 if (c == '\0')
648 c = '\n';
649 if (parse_add_char(self, c) < 0)
650 return -1;
651 self->state = IN_FIELD;
652 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000654 case IN_FIELD:
655 /* in unquoted field */
656 if (c == '\n' || c == '\r' || c == '\0') {
657 /* end of line - return [fields] */
658 if (parse_save_field(self) < 0)
659 return -1;
660 self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
661 }
662 else if (c == dialect->escapechar) {
663 /* possible escaped character */
664 self->state = ESCAPED_CHAR;
665 }
666 else if (c == dialect->delimiter) {
667 /* save field - wait for new field */
668 if (parse_save_field(self) < 0)
669 return -1;
670 self->state = START_FIELD;
671 }
672 else {
673 /* normal character - save in field */
674 if (parse_add_char(self, c) < 0)
675 return -1;
676 }
677 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 case IN_QUOTED_FIELD:
680 /* in quoted field */
681 if (c == '\0')
682 ;
683 else if (c == dialect->escapechar) {
684 /* Possible escape character */
685 self->state = ESCAPE_IN_QUOTED_FIELD;
686 }
687 else if (c == dialect->quotechar &&
688 dialect->quoting != QUOTE_NONE) {
689 if (dialect->doublequote) {
690 /* doublequote; " represented by "" */
691 self->state = QUOTE_IN_QUOTED_FIELD;
692 }
693 else {
694 /* end of quote part of field */
695 self->state = IN_FIELD;
696 }
697 }
698 else {
699 /* normal character - save in field */
700 if (parse_add_char(self, c) < 0)
701 return -1;
702 }
703 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000704
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000705 case ESCAPE_IN_QUOTED_FIELD:
706 if (c == '\0')
707 c = '\n';
708 if (parse_add_char(self, c) < 0)
709 return -1;
710 self->state = IN_QUOTED_FIELD;
711 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000712
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 case QUOTE_IN_QUOTED_FIELD:
714 /* doublequote - seen a quote in an quoted field */
715 if (dialect->quoting != QUOTE_NONE &&
716 c == dialect->quotechar) {
717 /* save "" as " */
718 if (parse_add_char(self, c) < 0)
719 return -1;
720 self->state = IN_QUOTED_FIELD;
721 }
722 else if (c == dialect->delimiter) {
723 /* save field - wait for new field */
724 if (parse_save_field(self) < 0)
725 return -1;
726 self->state = START_FIELD;
727 }
728 else if (c == '\n' || c == '\r' || c == '\0') {
729 /* end of line - return [fields] */
730 if (parse_save_field(self) < 0)
731 return -1;
732 self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
733 }
734 else if (!dialect->strict) {
735 if (parse_add_char(self, c) < 0)
736 return -1;
737 self->state = IN_FIELD;
738 }
739 else {
740 /* illegal */
Antoine Pitroue7672d32012-05-16 11:33:08 +0200741 PyErr_Format(_csvstate_global->error_obj, "'%c' expected after '%c'",
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000742 dialect->delimiter,
743 dialect->quotechar);
744 return -1;
745 }
746 break;
Skip Montanarob4a04172003-03-20 23:29:12 +0000747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000748 case EAT_CRNL:
749 if (c == '\n' || c == '\r')
750 ;
751 else if (c == '\0')
752 self->state = START_RECORD;
753 else {
Antoine Pitroue7672d32012-05-16 11:33:08 +0200754 PyErr_Format(_csvstate_global->error_obj, "new-line character seen in unquoted field - do you need to open the file in universal-newline mode?");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 return -1;
756 }
757 break;
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000759 }
760 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000761}
762
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000763static int
764parse_reset(ReaderObj *self)
765{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000766 Py_XDECREF(self->fields);
767 self->fields = PyList_New(0);
768 if (self->fields == NULL)
769 return -1;
770 self->field_len = 0;
771 self->state = START_RECORD;
772 self->numeric_field = 0;
773 return 0;
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000774}
Skip Montanarob4a04172003-03-20 23:29:12 +0000775
776static PyObject *
Skip Montanarob4a04172003-03-20 23:29:12 +0000777Reader_iternext(ReaderObj *self)
778{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000779 PyObject *fields = NULL;
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200780 Py_UCS4 c;
781 Py_ssize_t pos, linelen;
782 unsigned int kind;
783 void *data;
784 PyObject *lineobj;
Skip Montanarob4a04172003-03-20 23:29:12 +0000785
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000786 if (parse_reset(self) < 0)
787 return NULL;
788 do {
789 lineobj = PyIter_Next(self->input_iter);
790 if (lineobj == NULL) {
791 /* End of input OR exception */
Senthil Kumaran67b7b982012-09-25 02:30:27 -0700792 if (!PyErr_Occurred() && (self->field_len != 0 ||
793 self->state == IN_QUOTED_FIELD)) {
794 if (self->dialect->strict)
Senthil Kumaran49d13022012-09-25 02:37:20 -0700795 PyErr_SetString(_csvstate_global->error_obj,
796 "unexpected end of data");
Senthil Kumaran67b7b982012-09-25 02:30:27 -0700797 else if (parse_save_field(self) >= 0)
798 break;
799 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000800 return NULL;
801 }
802 if (!PyUnicode_Check(lineobj)) {
Antoine Pitroue7672d32012-05-16 11:33:08 +0200803 PyErr_Format(_csvstate_global->error_obj,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000804 "iterator should return strings, "
805 "not %.200s "
806 "(did you open the file in text mode?)",
807 lineobj->ob_type->tp_name
808 );
809 Py_DECREF(lineobj);
810 return NULL;
811 }
Stefan Krahe6996ed2012-11-02 14:44:20 +0100812 if (PyUnicode_READY(lineobj) == -1) {
813 Py_DECREF(lineobj);
814 return NULL;
815 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 ++self->line_num;
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200817 kind = PyUnicode_KIND(lineobj);
818 data = PyUnicode_DATA(lineobj);
819 pos = 0;
820 linelen = PyUnicode_GET_LENGTH(lineobj);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000821 while (linelen--) {
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200822 c = PyUnicode_READ(kind, data, pos);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000823 if (c == '\0') {
Skip Montanarob4a04172003-03-20 23:29:12 +0000824 Py_DECREF(lineobj);
Antoine Pitroue7672d32012-05-16 11:33:08 +0200825 PyErr_Format(_csvstate_global->error_obj,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000826 "line contains NULL byte");
827 goto err;
828 }
829 if (parse_process_char(self, c) < 0) {
830 Py_DECREF(lineobj);
831 goto err;
832 }
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200833 pos++;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000834 }
835 Py_DECREF(lineobj);
836 if (parse_process_char(self, 0) < 0)
837 goto err;
838 } while (self->state != START_RECORD);
Skip Montanarob4a04172003-03-20 23:29:12 +0000839
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000840 fields = self->fields;
841 self->fields = NULL;
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000842err:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000843 return fields;
Skip Montanarob4a04172003-03-20 23:29:12 +0000844}
845
846static void
847Reader_dealloc(ReaderObj *self)
848{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000849 PyObject_GC_UnTrack(self);
850 Py_XDECREF(self->dialect);
851 Py_XDECREF(self->input_iter);
852 Py_XDECREF(self->fields);
853 if (self->field != NULL)
854 PyMem_Free(self->field);
855 PyObject_GC_Del(self);
Jeremy Hylton42a8aed2003-04-14 02:20:55 +0000856}
857
858static int
859Reader_traverse(ReaderObj *self, visitproc visit, void *arg)
860{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 Py_VISIT(self->dialect);
862 Py_VISIT(self->input_iter);
863 Py_VISIT(self->fields);
864 return 0;
Jeremy Hylton42a8aed2003-04-14 02:20:55 +0000865}
866
867static int
868Reader_clear(ReaderObj *self)
869{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000870 Py_CLEAR(self->dialect);
871 Py_CLEAR(self->input_iter);
872 Py_CLEAR(self->fields);
873 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000874}
875
876PyDoc_STRVAR(Reader_Type_doc,
877"CSV reader\n"
878"\n"
879"Reader objects are responsible for reading and parsing tabular data\n"
880"in CSV format.\n"
881);
882
883static struct PyMethodDef Reader_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000884 { NULL, NULL }
Skip Montanarob4a04172003-03-20 23:29:12 +0000885};
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000886#define R_OFF(x) offsetof(ReaderObj, x)
887
888static struct PyMemberDef Reader_memberlist[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 { "dialect", T_OBJECT, R_OFF(dialect), READONLY },
890 { "line_num", T_ULONG, R_OFF(line_num), READONLY },
891 { NULL }
Andrew McNamaraf69d94f2005-01-13 11:30:54 +0000892};
893
Skip Montanarob4a04172003-03-20 23:29:12 +0000894
895static PyTypeObject Reader_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000896 PyVarObject_HEAD_INIT(NULL, 0)
897 "_csv.reader", /*tp_name*/
898 sizeof(ReaderObj), /*tp_basicsize*/
899 0, /*tp_itemsize*/
900 /* methods */
901 (destructor)Reader_dealloc, /*tp_dealloc*/
902 (printfunc)0, /*tp_print*/
903 (getattrfunc)0, /*tp_getattr*/
904 (setattrfunc)0, /*tp_setattr*/
905 0, /*tp_reserved*/
906 (reprfunc)0, /*tp_repr*/
907 0, /*tp_as_number*/
908 0, /*tp_as_sequence*/
909 0, /*tp_as_mapping*/
910 (hashfunc)0, /*tp_hash*/
911 (ternaryfunc)0, /*tp_call*/
912 (reprfunc)0, /*tp_str*/
913 0, /*tp_getattro*/
914 0, /*tp_setattro*/
915 0, /*tp_as_buffer*/
916 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
917 Py_TPFLAGS_HAVE_GC, /*tp_flags*/
918 Reader_Type_doc, /*tp_doc*/
919 (traverseproc)Reader_traverse, /*tp_traverse*/
920 (inquiry)Reader_clear, /*tp_clear*/
921 0, /*tp_richcompare*/
922 0, /*tp_weaklistoffset*/
923 PyObject_SelfIter, /*tp_iter*/
924 (getiterfunc)Reader_iternext, /*tp_iternext*/
925 Reader_methods, /*tp_methods*/
926 Reader_memberlist, /*tp_members*/
927 0, /*tp_getset*/
Skip Montanarob4a04172003-03-20 23:29:12 +0000928
929};
930
931static PyObject *
932csv_reader(PyObject *module, PyObject *args, PyObject *keyword_args)
933{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 PyObject * iterator, * dialect = NULL;
935 ReaderObj * self = PyObject_GC_New(ReaderObj, &Reader_Type);
Skip Montanarob4a04172003-03-20 23:29:12 +0000936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000937 if (!self)
938 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +0000939
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000940 self->dialect = NULL;
941 self->fields = NULL;
942 self->input_iter = NULL;
943 self->field = NULL;
944 self->field_size = 0;
945 self->line_num = 0;
Andrew McNamara0f0599d2005-01-12 09:45:18 +0000946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000947 if (parse_reset(self) < 0) {
948 Py_DECREF(self);
949 return NULL;
950 }
Skip Montanarob4a04172003-03-20 23:29:12 +0000951
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000952 if (!PyArg_UnpackTuple(args, "", 1, 2, &iterator, &dialect)) {
953 Py_DECREF(self);
954 return NULL;
955 }
956 self->input_iter = PyObject_GetIter(iterator);
957 if (self->input_iter == NULL) {
958 PyErr_SetString(PyExc_TypeError,
959 "argument 1 must be an iterator");
960 Py_DECREF(self);
961 return NULL;
962 }
963 self->dialect = (DialectObj *)_call_dialect(dialect, keyword_args);
964 if (self->dialect == NULL) {
965 Py_DECREF(self);
966 return NULL;
967 }
Skip Montanarob4a04172003-03-20 23:29:12 +0000968
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000969 PyObject_GC_Track(self);
970 return (PyObject *)self;
Skip Montanarob4a04172003-03-20 23:29:12 +0000971}
972
973/*
974 * WRITER
975 */
976/* ---------------------------------------------------------------- */
977static void
978join_reset(WriterObj *self)
979{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000980 self->rec_len = 0;
981 self->num_fields = 0;
Skip Montanarob4a04172003-03-20 23:29:12 +0000982}
983
984#define MEM_INCR 32768
985
986/* Calculate new record length or append field to record. Return new
987 * record length.
988 */
Antoine Pitrou40455752010-08-15 18:51:10 +0000989static Py_ssize_t
Antoine Pitrou77ea6402011-10-07 04:26:55 +0200990join_append_data(WriterObj *self, unsigned int field_kind, void *field_data,
991 Py_ssize_t field_len, int quote_empty, int *quoted,
992 int copy_phase)
Skip Montanarob4a04172003-03-20 23:29:12 +0000993{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000994 DialectObj *dialect = self->dialect;
995 int i;
Antoine Pitrou40455752010-08-15 18:51:10 +0000996 Py_ssize_t rec_len;
Andrew McNamarac89f2842005-01-12 07:44:42 +0000997
998#define ADDCH(c) \
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 do {\
1000 if (copy_phase) \
1001 self->rec[rec_len] = c;\
1002 rec_len++;\
1003 } while(0)
Andrew McNamarac89f2842005-01-12 07:44:42 +00001004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001005 rec_len = self->rec_len;
Skip Montanarob4a04172003-03-20 23:29:12 +00001006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 /* If this is not the first field we need a field separator */
1008 if (self->num_fields > 0)
1009 ADDCH(dialect->delimiter);
Andrew McNamarac89f2842005-01-12 07:44:42 +00001010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 /* Handle preceding quote */
1012 if (copy_phase && *quoted)
1013 ADDCH(dialect->quotechar);
Andrew McNamarac89f2842005-01-12 07:44:42 +00001014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001015 /* Copy/count field data */
1016 /* If field is null just pass over */
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001017 for (i = 0; field_data && (i < field_len); i++) {
1018 Py_UCS4 c = PyUnicode_READ(field_kind, field_data, i);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001019 int want_escape = 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001020
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001021 if (c == dialect->delimiter ||
1022 c == dialect->escapechar ||
1023 c == dialect->quotechar ||
Martin v. Löwis5f4f4c52011-11-01 18:42:23 +01001024 PyUnicode_FindChar(
1025 dialect->lineterminator, c, 0,
1026 PyUnicode_GET_LENGTH(dialect->lineterminator), 1) >= 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001027 if (dialect->quoting == QUOTE_NONE)
1028 want_escape = 1;
1029 else {
1030 if (c == dialect->quotechar) {
1031 if (dialect->doublequote)
1032 ADDCH(dialect->quotechar);
1033 else
1034 want_escape = 1;
1035 }
1036 if (!want_escape)
1037 *quoted = 1;
1038 }
1039 if (want_escape) {
1040 if (!dialect->escapechar) {
Antoine Pitroue7672d32012-05-16 11:33:08 +02001041 PyErr_Format(_csvstate_global->error_obj,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001042 "need to escape, but no escapechar set");
1043 return -1;
1044 }
1045 ADDCH(dialect->escapechar);
1046 }
1047 }
1048 /* Copy field character into record buffer.
1049 */
1050 ADDCH(c);
1051 }
Skip Montanarob4a04172003-03-20 23:29:12 +00001052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001053 /* If field is empty check if it needs to be quoted.
1054 */
1055 if (i == 0 && quote_empty) {
1056 if (dialect->quoting == QUOTE_NONE) {
Antoine Pitroue7672d32012-05-16 11:33:08 +02001057 PyErr_Format(_csvstate_global->error_obj,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001058 "single empty field record must be quoted");
1059 return -1;
1060 }
1061 else
1062 *quoted = 1;
1063 }
Skip Montanarob4a04172003-03-20 23:29:12 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 if (*quoted) {
1066 if (copy_phase)
1067 ADDCH(dialect->quotechar);
1068 else
1069 rec_len += 2;
1070 }
1071 return rec_len;
Andrew McNamarac89f2842005-01-12 07:44:42 +00001072#undef ADDCH
Skip Montanarob4a04172003-03-20 23:29:12 +00001073}
1074
1075static int
Antoine Pitrou40455752010-08-15 18:51:10 +00001076join_check_rec_size(WriterObj *self, Py_ssize_t rec_len)
Skip Montanarob4a04172003-03-20 23:29:12 +00001077{
Amaury Forgeot d'Arc9c74b142008-06-18 00:47:36 +00001078
Antoine Pitrou40455752010-08-15 18:51:10 +00001079 if (rec_len < 0 || rec_len > PY_SSIZE_T_MAX - MEM_INCR) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 PyErr_NoMemory();
1081 return 0;
1082 }
Amaury Forgeot d'Arc9c74b142008-06-18 00:47:36 +00001083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001084 if (rec_len > self->rec_size) {
1085 if (self->rec_size == 0) {
1086 self->rec_size = (rec_len / MEM_INCR + 1) * MEM_INCR;
1087 if (self->rec != NULL)
1088 PyMem_Free(self->rec);
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001089 self->rec = PyMem_New(Py_UCS4, self->rec_size);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 }
1091 else {
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001092 Py_UCS4* old_rec = self->rec;
Skip Montanarob4a04172003-03-20 23:29:12 +00001093
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001094 self->rec_size = (rec_len / MEM_INCR + 1) * MEM_INCR;
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001095 self->rec = PyMem_Resize(old_rec, Py_UCS4, self->rec_size);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001096 if (self->rec == NULL)
1097 PyMem_Free(old_rec);
1098 }
1099 if (self->rec == NULL) {
1100 PyErr_NoMemory();
1101 return 0;
1102 }
1103 }
1104 return 1;
Skip Montanarob4a04172003-03-20 23:29:12 +00001105}
1106
1107static int
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001108join_append(WriterObj *self, PyObject *field, int *quoted, int quote_empty)
Skip Montanarob4a04172003-03-20 23:29:12 +00001109{
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001110 unsigned int field_kind = -1;
1111 void *field_data = NULL;
1112 Py_ssize_t field_len = 0;
Antoine Pitrou40455752010-08-15 18:51:10 +00001113 Py_ssize_t rec_len;
Skip Montanarob4a04172003-03-20 23:29:12 +00001114
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001115 if (field != NULL) {
Stefan Krahe6996ed2012-11-02 14:44:20 +01001116 if (PyUnicode_READY(field) == -1)
1117 return 0;
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001118 field_kind = PyUnicode_KIND(field);
1119 field_data = PyUnicode_DATA(field);
1120 field_len = PyUnicode_GET_LENGTH(field);
1121 }
1122 rec_len = join_append_data(self, field_kind, field_data, field_len,
1123 quote_empty, quoted, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001124 if (rec_len < 0)
1125 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 /* grow record buffer if necessary */
1128 if (!join_check_rec_size(self, rec_len))
1129 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001130
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001131 self->rec_len = join_append_data(self, field_kind, field_data, field_len,
1132 quote_empty, quoted, 1);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001133 self->num_fields++;
Skip Montanarob4a04172003-03-20 23:29:12 +00001134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001135 return 1;
Skip Montanarob4a04172003-03-20 23:29:12 +00001136}
1137
1138static int
1139join_append_lineterminator(WriterObj *self)
1140{
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001141 Py_ssize_t terminator_len, i;
1142 unsigned int term_kind;
1143 void *term_data;
Skip Montanarob4a04172003-03-20 23:29:12 +00001144
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001145 terminator_len = PyUnicode_GET_LENGTH(self->dialect->lineterminator);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001146 if (terminator_len == -1)
1147 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001149 /* grow record buffer if necessary */
1150 if (!join_check_rec_size(self, self->rec_len + terminator_len))
1151 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001152
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001153 term_kind = PyUnicode_KIND(self->dialect->lineterminator);
1154 term_data = PyUnicode_DATA(self->dialect->lineterminator);
1155 for (i = 0; i < terminator_len; i++)
1156 self->rec[self->rec_len + i] = PyUnicode_READ(term_kind, term_data, i);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001157 self->rec_len += terminator_len;
Skip Montanarob4a04172003-03-20 23:29:12 +00001158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001159 return 1;
Skip Montanarob4a04172003-03-20 23:29:12 +00001160}
1161
1162PyDoc_STRVAR(csv_writerow_doc,
Skip Montanaro860fc0b2003-04-12 18:57:52 +00001163"writerow(sequence)\n"
Skip Montanarob4a04172003-03-20 23:29:12 +00001164"\n"
Skip Montanaro860fc0b2003-04-12 18:57:52 +00001165"Construct and write a CSV record from a sequence of fields. Non-string\n"
Skip Montanarob4a04172003-03-20 23:29:12 +00001166"elements will be converted to string.");
1167
1168static PyObject *
1169csv_writerow(WriterObj *self, PyObject *seq)
1170{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001171 DialectObj *dialect = self->dialect;
Antoine Pitrou40455752010-08-15 18:51:10 +00001172 Py_ssize_t len, i;
Antoine Pitrou72ca65d2011-10-07 04:35:30 +02001173 PyObject *line, *result;
Skip Montanarob4a04172003-03-20 23:29:12 +00001174
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001175 if (!PySequence_Check(seq))
Antoine Pitroue7672d32012-05-16 11:33:08 +02001176 return PyErr_Format(_csvstate_global->error_obj, "sequence expected");
Skip Montanarob4a04172003-03-20 23:29:12 +00001177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 len = PySequence_Length(seq);
1179 if (len < 0)
1180 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001181
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001182 /* Join all fields in internal buffer.
1183 */
1184 join_reset(self);
1185 for (i = 0; i < len; i++) {
1186 PyObject *field;
1187 int append_ok;
1188 int quoted;
Skip Montanarob4a04172003-03-20 23:29:12 +00001189
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001190 field = PySequence_GetItem(seq, i);
1191 if (field == NULL)
1192 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001193
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001194 switch (dialect->quoting) {
1195 case QUOTE_NONNUMERIC:
1196 quoted = !PyNumber_Check(field);
1197 break;
1198 case QUOTE_ALL:
1199 quoted = 1;
1200 break;
1201 default:
1202 quoted = 0;
1203 break;
1204 }
Skip Montanarob4a04172003-03-20 23:29:12 +00001205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 if (PyUnicode_Check(field)) {
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001207 append_ok = join_append(self, field, &quoted, len == 1);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001208 Py_DECREF(field);
1209 }
1210 else if (field == Py_None) {
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001211 append_ok = join_append(self, NULL, &quoted, len == 1);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 Py_DECREF(field);
1213 }
1214 else {
1215 PyObject *str;
Skip Montanarob4a04172003-03-20 23:29:12 +00001216
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001217 str = PyObject_Str(field);
1218 Py_DECREF(field);
1219 if (str == NULL)
1220 return NULL;
Antoine Pitrou77ea6402011-10-07 04:26:55 +02001221 append_ok = join_append(self, str, &quoted, len == 1);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 Py_DECREF(str);
1223 }
1224 if (!append_ok)
1225 return NULL;
1226 }
Skip Montanarob4a04172003-03-20 23:29:12 +00001227
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001228 /* Add line terminator.
1229 */
1230 if (!join_append_lineterminator(self))
1231 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001232
Antoine Pitrou72ca65d2011-10-07 04:35:30 +02001233 line = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND,
1234 (void *) self->rec, self->rec_len);
1235 if (line == NULL)
1236 return NULL;
1237 result = PyObject_CallFunctionObjArgs(self->writeline, line, NULL);
1238 Py_DECREF(line);
1239 return result;
Skip Montanarob4a04172003-03-20 23:29:12 +00001240}
1241
Skip Montanaro860fc0b2003-04-12 18:57:52 +00001242PyDoc_STRVAR(csv_writerows_doc,
1243"writerows(sequence of sequences)\n"
1244"\n"
1245"Construct and write a series of sequences to a csv file. Non-string\n"
1246"elements will be converted to string.");
1247
Skip Montanarob4a04172003-03-20 23:29:12 +00001248static PyObject *
1249csv_writerows(WriterObj *self, PyObject *seqseq)
1250{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001251 PyObject *row_iter, *row_obj, *result;
Skip Montanarob4a04172003-03-20 23:29:12 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 row_iter = PyObject_GetIter(seqseq);
1254 if (row_iter == NULL) {
1255 PyErr_SetString(PyExc_TypeError,
1256 "writerows() argument must be iterable");
1257 return NULL;
1258 }
1259 while ((row_obj = PyIter_Next(row_iter))) {
1260 result = csv_writerow(self, row_obj);
1261 Py_DECREF(row_obj);
1262 if (!result) {
1263 Py_DECREF(row_iter);
1264 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001265 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 else
1267 Py_DECREF(result);
1268 }
1269 Py_DECREF(row_iter);
1270 if (PyErr_Occurred())
1271 return NULL;
1272 Py_INCREF(Py_None);
1273 return Py_None;
Skip Montanarob4a04172003-03-20 23:29:12 +00001274}
1275
1276static struct PyMethodDef Writer_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001277 { "writerow", (PyCFunction)csv_writerow, METH_O, csv_writerow_doc},
1278 { "writerows", (PyCFunction)csv_writerows, METH_O, csv_writerows_doc},
1279 { NULL, NULL }
Skip Montanarob4a04172003-03-20 23:29:12 +00001280};
1281
1282#define W_OFF(x) offsetof(WriterObj, x)
1283
1284static struct PyMemberDef Writer_memberlist[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 { "dialect", T_OBJECT, W_OFF(dialect), READONLY },
1286 { NULL }
Skip Montanarob4a04172003-03-20 23:29:12 +00001287};
1288
1289static void
1290Writer_dealloc(WriterObj *self)
1291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001292 PyObject_GC_UnTrack(self);
1293 Py_XDECREF(self->dialect);
1294 Py_XDECREF(self->writeline);
1295 if (self->rec != NULL)
1296 PyMem_Free(self->rec);
1297 PyObject_GC_Del(self);
Jeremy Hylton42a8aed2003-04-14 02:20:55 +00001298}
1299
1300static int
1301Writer_traverse(WriterObj *self, visitproc visit, void *arg)
1302{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001303 Py_VISIT(self->dialect);
1304 Py_VISIT(self->writeline);
1305 return 0;
Jeremy Hylton42a8aed2003-04-14 02:20:55 +00001306}
1307
1308static int
1309Writer_clear(WriterObj *self)
1310{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 Py_CLEAR(self->dialect);
1312 Py_CLEAR(self->writeline);
1313 return 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001314}
1315
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001316PyDoc_STRVAR(Writer_Type_doc,
Skip Montanarob4a04172003-03-20 23:29:12 +00001317"CSV writer\n"
1318"\n"
1319"Writer objects are responsible for generating tabular data\n"
1320"in CSV format from sequence input.\n"
1321);
1322
1323static PyTypeObject Writer_Type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001324 PyVarObject_HEAD_INIT(NULL, 0)
1325 "_csv.writer", /*tp_name*/
1326 sizeof(WriterObj), /*tp_basicsize*/
1327 0, /*tp_itemsize*/
1328 /* methods */
1329 (destructor)Writer_dealloc, /*tp_dealloc*/
1330 (printfunc)0, /*tp_print*/
1331 (getattrfunc)0, /*tp_getattr*/
1332 (setattrfunc)0, /*tp_setattr*/
1333 0, /*tp_reserved*/
1334 (reprfunc)0, /*tp_repr*/
1335 0, /*tp_as_number*/
1336 0, /*tp_as_sequence*/
1337 0, /*tp_as_mapping*/
1338 (hashfunc)0, /*tp_hash*/
1339 (ternaryfunc)0, /*tp_call*/
1340 (reprfunc)0, /*tp_str*/
1341 0, /*tp_getattro*/
1342 0, /*tp_setattro*/
1343 0, /*tp_as_buffer*/
1344 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
1345 Py_TPFLAGS_HAVE_GC, /*tp_flags*/
1346 Writer_Type_doc,
1347 (traverseproc)Writer_traverse, /*tp_traverse*/
1348 (inquiry)Writer_clear, /*tp_clear*/
1349 0, /*tp_richcompare*/
1350 0, /*tp_weaklistoffset*/
1351 (getiterfunc)0, /*tp_iter*/
1352 (getiterfunc)0, /*tp_iternext*/
1353 Writer_methods, /*tp_methods*/
1354 Writer_memberlist, /*tp_members*/
1355 0, /*tp_getset*/
Skip Montanarob4a04172003-03-20 23:29:12 +00001356};
1357
1358static PyObject *
1359csv_writer(PyObject *module, PyObject *args, PyObject *keyword_args)
1360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 PyObject * output_file, * dialect = NULL;
1362 WriterObj * self = PyObject_GC_New(WriterObj, &Writer_Type);
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001363 _Py_IDENTIFIER(write);
Skip Montanarob4a04172003-03-20 23:29:12 +00001364
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001365 if (!self)
1366 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001368 self->dialect = NULL;
1369 self->writeline = NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001371 self->rec = NULL;
1372 self->rec_size = 0;
1373 self->rec_len = 0;
1374 self->num_fields = 0;
Skip Montanarob4a04172003-03-20 23:29:12 +00001375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 if (!PyArg_UnpackTuple(args, "", 1, 2, &output_file, &dialect)) {
1377 Py_DECREF(self);
1378 return NULL;
1379 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001380 self->writeline = _PyObject_GetAttrId(output_file, &PyId_write);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001381 if (self->writeline == NULL || !PyCallable_Check(self->writeline)) {
1382 PyErr_SetString(PyExc_TypeError,
1383 "argument 1 must have a \"write\" method");
1384 Py_DECREF(self);
1385 return NULL;
1386 }
1387 self->dialect = (DialectObj *)_call_dialect(dialect, keyword_args);
1388 if (self->dialect == NULL) {
1389 Py_DECREF(self);
1390 return NULL;
1391 }
1392 PyObject_GC_Track(self);
1393 return (PyObject *)self;
Skip Montanarob4a04172003-03-20 23:29:12 +00001394}
1395
1396/*
1397 * DIALECT REGISTRY
1398 */
1399static PyObject *
1400csv_list_dialects(PyObject *module, PyObject *args)
1401{
Antoine Pitroue7672d32012-05-16 11:33:08 +02001402 return PyDict_Keys(_csvstate_global->dialects);
Skip Montanarob4a04172003-03-20 23:29:12 +00001403}
1404
1405static PyObject *
Andrew McNamara86625972005-01-11 01:28:33 +00001406csv_register_dialect(PyObject *module, PyObject *args, PyObject *kwargs)
Skip Montanarob4a04172003-03-20 23:29:12 +00001407{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001408 PyObject *name_obj, *dialect_obj = NULL;
1409 PyObject *dialect;
Skip Montanarob4a04172003-03-20 23:29:12 +00001410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001411 if (!PyArg_UnpackTuple(args, "", 1, 2, &name_obj, &dialect_obj))
1412 return NULL;
Stefan Krahe6996ed2012-11-02 14:44:20 +01001413 if (!PyUnicode_Check(name_obj)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001414 PyErr_SetString(PyExc_TypeError,
Stefan Krahe6996ed2012-11-02 14:44:20 +01001415 "dialect name must be a string");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 return NULL;
1417 }
Stefan Krahe6996ed2012-11-02 14:44:20 +01001418 if (PyUnicode_READY(name_obj) == -1)
1419 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001420 dialect = _call_dialect(dialect_obj, kwargs);
1421 if (dialect == NULL)
1422 return NULL;
Antoine Pitroue7672d32012-05-16 11:33:08 +02001423 if (PyDict_SetItem(_csvstate_global->dialects, name_obj, dialect) < 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001424 Py_DECREF(dialect);
1425 return NULL;
1426 }
1427 Py_DECREF(dialect);
1428 Py_INCREF(Py_None);
1429 return Py_None;
Skip Montanarob4a04172003-03-20 23:29:12 +00001430}
1431
1432static PyObject *
Skip Montanaro577c7a72003-04-12 19:17:14 +00001433csv_unregister_dialect(PyObject *module, PyObject *name_obj)
Skip Montanarob4a04172003-03-20 23:29:12 +00001434{
Antoine Pitroue7672d32012-05-16 11:33:08 +02001435 if (PyDict_DelItem(_csvstate_global->dialects, name_obj) < 0)
1436 return PyErr_Format(_csvstate_global->error_obj, "unknown dialect");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 Py_INCREF(Py_None);
1438 return Py_None;
Skip Montanarob4a04172003-03-20 23:29:12 +00001439}
1440
1441static PyObject *
Skip Montanaro577c7a72003-04-12 19:17:14 +00001442csv_get_dialect(PyObject *module, PyObject *name_obj)
Skip Montanarob4a04172003-03-20 23:29:12 +00001443{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001444 return get_dialect_from_registry(name_obj);
Skip Montanarob4a04172003-03-20 23:29:12 +00001445}
1446
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001447static PyObject *
Andrew McNamara31d88962005-01-12 03:45:10 +00001448csv_field_size_limit(PyObject *module, PyObject *args)
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001449{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001450 PyObject *new_limit = NULL;
Antoine Pitroue7672d32012-05-16 11:33:08 +02001451 long old_limit = _csvstate_global->field_limit;
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001452
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001453 if (!PyArg_UnpackTuple(args, "field_size_limit", 0, 1, &new_limit))
1454 return NULL;
1455 if (new_limit != NULL) {
1456 if (!PyLong_CheckExact(new_limit)) {
1457 PyErr_Format(PyExc_TypeError,
1458 "limit must be an integer");
1459 return NULL;
1460 }
Antoine Pitroue7672d32012-05-16 11:33:08 +02001461 _csvstate_global->field_limit = PyLong_AsLong(new_limit);
1462 if (_csvstate_global->field_limit == -1 && PyErr_Occurred()) {
1463 _csvstate_global->field_limit = old_limit;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001464 return NULL;
1465 }
1466 }
1467 return PyLong_FromLong(old_limit);
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001468}
1469
Skip Montanarob4a04172003-03-20 23:29:12 +00001470/*
1471 * MODULE
1472 */
1473
1474PyDoc_STRVAR(csv_module_doc,
1475"CSV parsing and writing.\n"
1476"\n"
1477"This module provides classes that assist in the reading and writing\n"
1478"of Comma Separated Value (CSV) files, and implements the interface\n"
1479"described by PEP 305. Although many CSV files are simple to parse,\n"
1480"the format is not formally defined by a stable specification and\n"
1481"is subtle enough that parsing lines of a CSV file with something\n"
1482"like line.split(\",\") is bound to fail. The module supports three\n"
1483"basic APIs: reading, writing, and registration of dialects.\n"
1484"\n"
1485"\n"
1486"DIALECT REGISTRATION:\n"
1487"\n"
1488"Readers and writers support a dialect argument, which is a convenient\n"
1489"handle on a group of settings. When the dialect argument is a string,\n"
1490"it identifies one of the dialects previously registered with the module.\n"
1491"If it is a class or instance, the attributes of the argument are used as\n"
1492"the settings for the reader or writer:\n"
1493"\n"
1494" class excel:\n"
1495" delimiter = ','\n"
1496" quotechar = '\"'\n"
1497" escapechar = None\n"
1498" doublequote = True\n"
1499" skipinitialspace = False\n"
Johannes Gijsbers8d3b9dd2004-08-15 12:23:10 +00001500" lineterminator = '\\r\\n'\n"
Skip Montanarob4a04172003-03-20 23:29:12 +00001501" quoting = QUOTE_MINIMAL\n"
1502"\n"
1503"SETTINGS:\n"
1504"\n"
1505" * quotechar - specifies a one-character string to use as the \n"
1506" quoting character. It defaults to '\"'.\n"
1507" * delimiter - specifies a one-character string to use as the \n"
1508" field separator. It defaults to ','.\n"
1509" * skipinitialspace - specifies how to interpret whitespace which\n"
1510" immediately follows a delimiter. It defaults to False, which\n"
1511" means that whitespace immediately following a delimiter is part\n"
1512" of the following field.\n"
1513" * lineterminator - specifies the character sequence which should \n"
1514" terminate rows.\n"
1515" * quoting - controls when quotes should be generated by the writer.\n"
1516" It can take on any of the following module constants:\n"
1517"\n"
1518" csv.QUOTE_MINIMAL means only when required, for example, when a\n"
1519" field contains either the quotechar or the delimiter\n"
1520" csv.QUOTE_ALL means that quotes are always placed around fields.\n"
1521" csv.QUOTE_NONNUMERIC means that quotes are always placed around\n"
Skip Montanaro148eb6a2003-12-02 18:57:47 +00001522" fields which do not parse as integers or floating point\n"
1523" numbers.\n"
Skip Montanarob4a04172003-03-20 23:29:12 +00001524" csv.QUOTE_NONE means that quotes are never placed around fields.\n"
1525" * escapechar - specifies a one-character string used to escape \n"
1526" the delimiter when quoting is set to QUOTE_NONE.\n"
1527" * doublequote - controls the handling of quotes inside fields. When\n"
1528" True, two consecutive quotes are interpreted as one during read,\n"
1529" and when writing, each quote character embedded in the data is\n"
1530" written as two quotes\n");
1531
1532PyDoc_STRVAR(csv_reader_doc,
1533" csv_reader = reader(iterable [, dialect='excel']\n"
1534" [optional keyword args])\n"
1535" for row in csv_reader:\n"
1536" process(row)\n"
1537"\n"
1538"The \"iterable\" argument can be any object that returns a line\n"
1539"of input for each iteration, such as a file object or a list. The\n"
1540"optional \"dialect\" parameter is discussed below. The function\n"
1541"also accepts optional keyword arguments which override settings\n"
1542"provided by the dialect.\n"
1543"\n"
1544"The returned object is an iterator. Each iteration returns a row\n"
Johannes Gijsbers8d3b9dd2004-08-15 12:23:10 +00001545"of the CSV file (which can span multiple input lines):\n");
Skip Montanarob4a04172003-03-20 23:29:12 +00001546
1547PyDoc_STRVAR(csv_writer_doc,
1548" csv_writer = csv.writer(fileobj [, dialect='excel']\n"
1549" [optional keyword args])\n"
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00001550" for row in sequence:\n"
Skip Montanarob4a04172003-03-20 23:29:12 +00001551" csv_writer.writerow(row)\n"
1552"\n"
1553" [or]\n"
1554"\n"
1555" csv_writer = csv.writer(fileobj [, dialect='excel']\n"
1556" [optional keyword args])\n"
1557" csv_writer.writerows(rows)\n"
1558"\n"
1559"The \"fileobj\" argument can be any object that supports the file API.\n");
1560
1561PyDoc_STRVAR(csv_list_dialects_doc,
1562"Return a list of all know dialect names.\n"
1563" names = csv.list_dialects()");
1564
1565PyDoc_STRVAR(csv_get_dialect_doc,
1566"Return the dialect instance associated with name.\n"
1567" dialect = csv.get_dialect(name)");
1568
1569PyDoc_STRVAR(csv_register_dialect_doc,
1570"Create a mapping from a string name to a dialect class.\n"
1571" dialect = csv.register_dialect(name, dialect)");
1572
1573PyDoc_STRVAR(csv_unregister_dialect_doc,
1574"Delete the name/dialect mapping associated with a string name.\n"
1575" csv.unregister_dialect(name)");
1576
Andrew McNamara31d88962005-01-12 03:45:10 +00001577PyDoc_STRVAR(csv_field_size_limit_doc,
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001578"Sets an upper limit on parsed fields.\n"
Andrew McNamara31d88962005-01-12 03:45:10 +00001579" csv.field_size_limit([limit])\n"
Andrew McNamarae4d05c42005-01-11 07:32:02 +00001580"\n"
1581"Returns old limit. If limit is not given, no new limit is set and\n"
1582"the old limit is returned");
1583
Skip Montanarob4a04172003-03-20 23:29:12 +00001584static struct PyMethodDef csv_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001585 { "reader", (PyCFunction)csv_reader,
1586 METH_VARARGS | METH_KEYWORDS, csv_reader_doc},
1587 { "writer", (PyCFunction)csv_writer,
1588 METH_VARARGS | METH_KEYWORDS, csv_writer_doc},
1589 { "list_dialects", (PyCFunction)csv_list_dialects,
1590 METH_NOARGS, csv_list_dialects_doc},
1591 { "register_dialect", (PyCFunction)csv_register_dialect,
1592 METH_VARARGS | METH_KEYWORDS, csv_register_dialect_doc},
1593 { "unregister_dialect", (PyCFunction)csv_unregister_dialect,
1594 METH_O, csv_unregister_dialect_doc},
1595 { "get_dialect", (PyCFunction)csv_get_dialect,
1596 METH_O, csv_get_dialect_doc},
1597 { "field_size_limit", (PyCFunction)csv_field_size_limit,
1598 METH_VARARGS, csv_field_size_limit_doc},
1599 { NULL, NULL }
Skip Montanarob4a04172003-03-20 23:29:12 +00001600};
1601
Martin v. Löwis1a214512008-06-11 05:26:20 +00001602static struct PyModuleDef _csvmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001603 PyModuleDef_HEAD_INIT,
1604 "_csv",
1605 csv_module_doc,
Antoine Pitroue7672d32012-05-16 11:33:08 +02001606 sizeof(_csvstate),
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001607 csv_methods,
1608 NULL,
Antoine Pitroue7672d32012-05-16 11:33:08 +02001609 _csv_traverse,
1610 _csv_clear,
1611 _csv_free
Martin v. Löwis1a214512008-06-11 05:26:20 +00001612};
1613
Skip Montanarob4a04172003-03-20 23:29:12 +00001614PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00001615PyInit__csv(void)
Skip Montanarob4a04172003-03-20 23:29:12 +00001616{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 PyObject *module;
1618 StyleDesc *style;
Skip Montanarob4a04172003-03-20 23:29:12 +00001619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001620 if (PyType_Ready(&Dialect_Type) < 0)
1621 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001623 if (PyType_Ready(&Reader_Type) < 0)
1624 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 if (PyType_Ready(&Writer_Type) < 0)
1627 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001629 /* Create the module and add the functions */
1630 module = PyModule_Create(&_csvmodule);
1631 if (module == NULL)
1632 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 /* Add version to the module. */
1635 if (PyModule_AddStringConstant(module, "__version__",
1636 MODULE_VERSION) == -1)
1637 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001638
Antoine Pitroue7672d32012-05-16 11:33:08 +02001639 /* Set the field limit */
1640 _csvstate(module)->field_limit = 128 * 1024;
1641 /* Do I still need to add this var to the Module Dict? */
1642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001643 /* Add _dialects dictionary */
Antoine Pitroue7672d32012-05-16 11:33:08 +02001644 _csvstate(module)->dialects = PyDict_New();
1645 if (_csvstate(module)->dialects == NULL)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001646 return NULL;
Antoine Pitroue7672d32012-05-16 11:33:08 +02001647 Py_INCREF(_csvstate(module)->dialects);
1648 if (PyModule_AddObject(module, "_dialects", _csvstate(module)->dialects))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 /* Add quote styles into dictionary */
1652 for (style = quote_styles; style->name; style++) {
1653 if (PyModule_AddIntConstant(module, style->name,
1654 style->style) == -1)
1655 return NULL;
1656 }
Skip Montanarob4a04172003-03-20 23:29:12 +00001657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001658 /* Add the Dialect type */
1659 Py_INCREF(&Dialect_Type);
1660 if (PyModule_AddObject(module, "Dialect", (PyObject *)&Dialect_Type))
1661 return NULL;
Skip Montanarob4a04172003-03-20 23:29:12 +00001662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001663 /* Add the CSV exception object to the module. */
Antoine Pitroue7672d32012-05-16 11:33:08 +02001664 _csvstate(module)->error_obj = PyErr_NewException("_csv.Error", NULL, NULL);
1665 if (_csvstate(module)->error_obj == NULL)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001666 return NULL;
Antoine Pitroue7672d32012-05-16 11:33:08 +02001667 Py_INCREF(_csvstate(module)->error_obj);
1668 PyModule_AddObject(module, "Error", _csvstate(module)->error_obj);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 return module;
Skip Montanarob4a04172003-03-20 23:29:12 +00001670}