fastcsv.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266
  1. import csv
  2. import cffi
  3. # IN-PROGRESS. See the demo at the end of the file
  4. def _make_ffi_from_dialect(dialect_name):
  5. dialect = csv.get_dialect(dialect_name)
  6. ffi = cffi.FFI()
  7. ffi.cdef("""
  8. long parse_line(char *rawline, long inputlength);
  9. """)
  10. d = {'quotechar': ord(dialect.quotechar),
  11. 'quoting': int(dialect.quoting),
  12. 'skipinitialspace': int(dialect.skipinitialspace),
  13. 'delimiter': ord(dialect.delimiter),
  14. 'doublequote': int(dialect.doublequote),
  15. 'strict': int(dialect.strict),
  16. }
  17. if dialect.escapechar is not None:
  18. d['is_escape_char'] = '== %d' % ord(dialect.escapechar)
  19. else:
  20. d['is_escape_char'] = '&& 0'
  21. ffi.set_source('_fastcsv_' + dialect_name, r'''
  22. typedef enum {
  23. START_RECORD, START_FIELD, ESCAPED_CHAR, IN_FIELD,
  24. IN_QUOTED_FIELD, ESCAPE_IN_QUOTED_FIELD, QUOTE_IN_QUOTED_FIELD,
  25. EAT_CRNL
  26. } ParserState;
  27. typedef enum {
  28. QUOTE_MINIMAL, QUOTE_ALL, QUOTE_NONNUMERIC, QUOTE_NONE
  29. } QuoteStyle;
  30. typedef struct {
  31. ParserState state; /* current CSV parse state */
  32. char *field; /* build current field in here */
  33. int field_size; /* size of allocated buffer */
  34. int field_len; /* length of current field */
  35. int numeric_field; /* treat field as numeric */
  36. } ReaderObj;
  37. static void
  38. parse_add_char(ReaderObj *self, char c)
  39. {
  40. *self->field++ = c;
  41. }
  42. static void
  43. parse_save_field(ReaderObj *self)
  44. {
  45. *self->field++ = 0;
  46. }
  47. static int
  48. parse_process_char(ReaderObj *self, char c)
  49. {
  50. switch (self->state) {
  51. case START_RECORD:
  52. /* start of record */
  53. if (c == '\0')
  54. /* empty line - return [] */
  55. break;
  56. else if (c == '\n' || c == '\r') {
  57. self->state = EAT_CRNL;
  58. break;
  59. }
  60. /* normal character - handle as START_FIELD */
  61. self->state = START_FIELD;
  62. /* fallthru */
  63. case START_FIELD:
  64. /* expecting field */
  65. if (c == '\n' || c == '\r' || c == '\0') {
  66. /* save empty field - return [fields] */
  67. parse_save_field(self);
  68. self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
  69. }
  70. else if (c == %(quotechar)d &&
  71. %(quoting)d != QUOTE_NONE) {
  72. /* start quoted field */
  73. self->state = IN_QUOTED_FIELD;
  74. }
  75. else if (c %(is_escape_char)s) {
  76. /* possible escaped character */
  77. self->state = ESCAPED_CHAR;
  78. }
  79. else if (c == ' ' && %(skipinitialspace)d)
  80. /* ignore space at start of field */
  81. ;
  82. else if (c == %(delimiter)d) {
  83. /* save empty field */
  84. parse_save_field(self);
  85. }
  86. else {
  87. /* begin new unquoted field */
  88. if (%(quoting)d == QUOTE_NONNUMERIC)
  89. self->numeric_field = 1;
  90. parse_add_char(self, c);
  91. self->state = IN_FIELD;
  92. }
  93. break;
  94. case ESCAPED_CHAR:
  95. if (c == '\0')
  96. c = '\n';
  97. parse_add_char(self, c);
  98. self->state = IN_FIELD;
  99. break;
  100. case IN_FIELD:
  101. /* in unquoted field */
  102. if (c == '\n' || c == '\r' || c == '\0') {
  103. /* end of line - return [fields] */
  104. parse_save_field(self);
  105. self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
  106. }
  107. else if (c %(is_escape_char)s) {
  108. /* possible escaped character */
  109. self->state = ESCAPED_CHAR;
  110. }
  111. else if (c == %(delimiter)d) {
  112. /* save field - wait for new field */
  113. parse_save_field(self);
  114. self->state = START_FIELD;
  115. }
  116. else {
  117. /* normal character - save in field */
  118. parse_add_char(self, c);
  119. }
  120. break;
  121. case IN_QUOTED_FIELD:
  122. /* in quoted field */
  123. if (c == '\0')
  124. ;
  125. else if (c %(is_escape_char)s) {
  126. /* Possible escape character */
  127. self->state = ESCAPE_IN_QUOTED_FIELD;
  128. }
  129. else if (c == %(quotechar)d &&
  130. %(quoting)d != QUOTE_NONE) {
  131. if (%(doublequote)d) {
  132. /* doublequote; " represented by "" */
  133. self->state = QUOTE_IN_QUOTED_FIELD;
  134. }
  135. else {
  136. /* end of quote part of field */
  137. self->state = IN_FIELD;
  138. }
  139. }
  140. else {
  141. /* normal character - save in field */
  142. parse_add_char(self, c);
  143. }
  144. break;
  145. case ESCAPE_IN_QUOTED_FIELD:
  146. if (c == '\0')
  147. c = '\n';
  148. parse_add_char(self, c);
  149. self->state = IN_QUOTED_FIELD;
  150. break;
  151. case QUOTE_IN_QUOTED_FIELD:
  152. /* doublequote - seen a quote in an quoted field */
  153. if (%(quoting)d != QUOTE_NONE &&
  154. c == %(quotechar)d) {
  155. /* save "" as " */
  156. parse_add_char(self, c);
  157. self->state = IN_QUOTED_FIELD;
  158. }
  159. else if (c == %(delimiter)d) {
  160. /* save field - wait for new field */
  161. parse_save_field(self);
  162. self->state = START_FIELD;
  163. }
  164. else if (c == '\n' || c == '\r' || c == '\0') {
  165. /* end of line - return [fields] */
  166. parse_save_field(self);
  167. self->state = (c == '\0' ? START_RECORD : EAT_CRNL);
  168. }
  169. else if (!%(strict)d) {
  170. parse_add_char(self, c);
  171. self->state = IN_FIELD;
  172. }
  173. else {
  174. /* illegal */
  175. /*PyErr_Format(error_obj, "'%%c' expected after '%%c'",
  176. dialect->delimiter,
  177. dialect->quotechar);*/
  178. return -1;
  179. }
  180. break;
  181. case EAT_CRNL:
  182. if (c == '\n' || c == '\r')
  183. ;
  184. else if (c == '\0')
  185. self->state = START_RECORD;
  186. else {
  187. /*PyErr_Format(error_obj, "new-line character seen in unquoted field - do you need to open the file in universal-newline mode?");*/
  188. return -1;
  189. }
  190. break;
  191. }
  192. return 0;
  193. }
  194. static void
  195. parse_reset(ReaderObj *self, char *rawline)
  196. {
  197. self->field = rawline;
  198. self->state = START_RECORD;
  199. self->numeric_field = 0;
  200. }
  201. long parse_line(char *rawline, long inputlength)
  202. {
  203. char *p;
  204. ReaderObj reader;
  205. parse_reset(&reader, rawline);
  206. for (p=rawline; inputlength > 0; inputlength--, p++) {
  207. if (parse_process_char(&reader, *p) < 0)
  208. return -1;
  209. }
  210. if (parse_process_char(&reader, 0) < 0)
  211. return -1;
  212. return reader.field - rawline - 1;
  213. }
  214. ''' % d)
  215. ffi.compile()
  216. def fastcsv_reader(f, dialect_name):
  217. try:
  218. module = __import__('_fastcsv_' + dialect_name)
  219. except ImportError:
  220. _make_ffi_from_dialect(dialect_name)
  221. module = __import__('_fastcsv_' + dialect_name)
  222. ffi, lib = module.ffi, module.lib
  223. #
  224. linelen = -1
  225. for line in f:
  226. if linelen <= len(line):
  227. linelen = 2 * len(line)
  228. rawline = ffi.new("char[]", linelen)
  229. ffi.buffer(rawline, len(line))[:] = line
  230. n = lib.parse_line(rawline, len(line))
  231. assert n >= 0
  232. yield ffi.buffer(rawline, n)[:].split('\x00')
  233. if __name__ == '__main__':
  234. csv.register_dialect('unixpwd', delimiter=':', quoting=csv.QUOTE_NONE)
  235. with open('/etc/passwd', 'rb') as f:
  236. reader = fastcsv_reader(f, 'unixpwd')
  237. for row in reader:
  238. print row