csv.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451
  1. """
  2. csv.py - read/write/investigate CSV files
  3. """
  4. import re
  5. import types
  6. from _csv import Error, __version__, writer, reader, register_dialect, \
  7. unregister_dialect, get_dialect, list_dialects, \
  8. field_size_limit, \
  9. QUOTE_MINIMAL, QUOTE_ALL, QUOTE_NONNUMERIC, QUOTE_NONE, \
  10. QUOTE_STRINGS, QUOTE_NOTNULL, \
  11. __doc__
  12. from _csv import Dialect as _Dialect
  13. from io import StringIO
  14. __all__ = ["QUOTE_MINIMAL", "QUOTE_ALL", "QUOTE_NONNUMERIC", "QUOTE_NONE",
  15. "QUOTE_STRINGS", "QUOTE_NOTNULL",
  16. "Error", "Dialect", "__doc__", "excel", "excel_tab",
  17. "field_size_limit", "reader", "writer",
  18. "register_dialect", "get_dialect", "list_dialects", "Sniffer",
  19. "unregister_dialect", "__version__", "DictReader", "DictWriter",
  20. "unix_dialect"]
  21. class Dialect:
  22. """Describe a CSV dialect.
  23. This must be subclassed (see csv.excel). Valid attributes are:
  24. delimiter, quotechar, escapechar, doublequote, skipinitialspace,
  25. lineterminator, quoting.
  26. """
  27. _name = ""
  28. _valid = False
  29. # placeholders
  30. delimiter = None
  31. quotechar = None
  32. escapechar = None
  33. doublequote = None
  34. skipinitialspace = None
  35. lineterminator = None
  36. quoting = None
  37. def __init__(self):
  38. if self.__class__ != Dialect:
  39. self._valid = True
  40. self._validate()
  41. def _validate(self):
  42. try:
  43. _Dialect(self)
  44. except TypeError as e:
  45. # We do this for compatibility with py2.3
  46. raise Error(str(e))
  47. class excel(Dialect):
  48. """Describe the usual properties of Excel-generated CSV files."""
  49. delimiter = ','
  50. quotechar = '"'
  51. doublequote = True
  52. skipinitialspace = False
  53. lineterminator = '\r\n'
  54. quoting = QUOTE_MINIMAL
  55. register_dialect("excel", excel)
  56. class excel_tab(excel):
  57. """Describe the usual properties of Excel-generated TAB-delimited files."""
  58. delimiter = '\t'
  59. register_dialect("excel-tab", excel_tab)
  60. class unix_dialect(Dialect):
  61. """Describe the usual properties of Unix-generated CSV files."""
  62. delimiter = ','
  63. quotechar = '"'
  64. doublequote = True
  65. skipinitialspace = False
  66. lineterminator = '\n'
  67. quoting = QUOTE_ALL
  68. register_dialect("unix", unix_dialect)
  69. class DictReader:
  70. def __init__(self, f, fieldnames=None, restkey=None, restval=None,
  71. dialect="excel", *args, **kwds):
  72. if fieldnames is not None and iter(fieldnames) is fieldnames:
  73. fieldnames = list(fieldnames)
  74. self._fieldnames = fieldnames # list of keys for the dict
  75. self.restkey = restkey # key to catch long rows
  76. self.restval = restval # default value for short rows
  77. self.reader = reader(f, dialect, *args, **kwds)
  78. self.dialect = dialect
  79. self.line_num = 0
  80. def __iter__(self):
  81. return self
  82. @property
  83. def fieldnames(self):
  84. if self._fieldnames is None:
  85. try:
  86. self._fieldnames = next(self.reader)
  87. except StopIteration:
  88. pass
  89. self.line_num = self.reader.line_num
  90. return self._fieldnames
  91. @fieldnames.setter
  92. def fieldnames(self, value):
  93. self._fieldnames = value
  94. def __next__(self):
  95. if self.line_num == 0:
  96. # Used only for its side effect.
  97. self.fieldnames
  98. row = next(self.reader)
  99. self.line_num = self.reader.line_num
  100. # unlike the basic reader, we prefer not to return blanks,
  101. # because we will typically wind up with a dict full of None
  102. # values
  103. while row == []:
  104. row = next(self.reader)
  105. d = dict(zip(self.fieldnames, row))
  106. lf = len(self.fieldnames)
  107. lr = len(row)
  108. if lf < lr:
  109. d[self.restkey] = row[lf:]
  110. elif lf > lr:
  111. for key in self.fieldnames[lr:]:
  112. d[key] = self.restval
  113. return d
  114. __class_getitem__ = classmethod(types.GenericAlias)
  115. class DictWriter:
  116. def __init__(self, f, fieldnames, restval="", extrasaction="raise",
  117. dialect="excel", *args, **kwds):
  118. if fieldnames is not None and iter(fieldnames) is fieldnames:
  119. fieldnames = list(fieldnames)
  120. self.fieldnames = fieldnames # list of keys for the dict
  121. self.restval = restval # for writing short dicts
  122. extrasaction = extrasaction.lower()
  123. if extrasaction not in ("raise", "ignore"):
  124. raise ValueError("extrasaction (%s) must be 'raise' or 'ignore'"
  125. % extrasaction)
  126. self.extrasaction = extrasaction
  127. self.writer = writer(f, dialect, *args, **kwds)
  128. def writeheader(self):
  129. header = dict(zip(self.fieldnames, self.fieldnames))
  130. return self.writerow(header)
  131. def _dict_to_list(self, rowdict):
  132. if self.extrasaction == "raise":
  133. wrong_fields = rowdict.keys() - self.fieldnames
  134. if wrong_fields:
  135. raise ValueError("dict contains fields not in fieldnames: "
  136. + ", ".join([repr(x) for x in wrong_fields]))
  137. return (rowdict.get(key, self.restval) for key in self.fieldnames)
  138. def writerow(self, rowdict):
  139. return self.writer.writerow(self._dict_to_list(rowdict))
  140. def writerows(self, rowdicts):
  141. return self.writer.writerows(map(self._dict_to_list, rowdicts))
  142. __class_getitem__ = classmethod(types.GenericAlias)
  143. class Sniffer:
  144. '''
  145. "Sniffs" the format of a CSV file (i.e. delimiter, quotechar)
  146. Returns a Dialect object.
  147. '''
  148. def __init__(self):
  149. # in case there is more than one possible delimiter
  150. self.preferred = [',', '\t', ';', ' ', ':']
  151. def sniff(self, sample, delimiters=None):
  152. """
  153. Returns a dialect (or None) corresponding to the sample
  154. """
  155. quotechar, doublequote, delimiter, skipinitialspace = \
  156. self._guess_quote_and_delimiter(sample, delimiters)
  157. if not delimiter:
  158. delimiter, skipinitialspace = self._guess_delimiter(sample,
  159. delimiters)
  160. if not delimiter:
  161. raise Error("Could not determine delimiter")
  162. class dialect(Dialect):
  163. _name = "sniffed"
  164. lineterminator = '\r\n'
  165. quoting = QUOTE_MINIMAL
  166. # escapechar = ''
  167. dialect.doublequote = doublequote
  168. dialect.delimiter = delimiter
  169. # _csv.reader won't accept a quotechar of ''
  170. dialect.quotechar = quotechar or '"'
  171. dialect.skipinitialspace = skipinitialspace
  172. return dialect
  173. def _guess_quote_and_delimiter(self, data, delimiters):
  174. """
  175. Looks for text enclosed between two identical quotes
  176. (the probable quotechar) which are preceded and followed
  177. by the same character (the probable delimiter).
  178. For example:
  179. ,'some text',
  180. The quote with the most wins, same with the delimiter.
  181. If there is no quotechar the delimiter can't be determined
  182. this way.
  183. """
  184. matches = []
  185. for restr in (r'(?P<delim>[^\w\n"\'])(?P<space> ?)(?P<quote>["\']).*?(?P=quote)(?P=delim)', # ,".*?",
  186. r'(?:^|\n)(?P<quote>["\']).*?(?P=quote)(?P<delim>[^\w\n"\'])(?P<space> ?)', # ".*?",
  187. r'(?P<delim>[^\w\n"\'])(?P<space> ?)(?P<quote>["\']).*?(?P=quote)(?:$|\n)', # ,".*?"
  188. r'(?:^|\n)(?P<quote>["\']).*?(?P=quote)(?:$|\n)'): # ".*?" (no delim, no space)
  189. regexp = re.compile(restr, re.DOTALL | re.MULTILINE)
  190. matches = regexp.findall(data)
  191. if matches:
  192. break
  193. if not matches:
  194. # (quotechar, doublequote, delimiter, skipinitialspace)
  195. return ('', False, None, 0)
  196. quotes = {}
  197. delims = {}
  198. spaces = 0
  199. groupindex = regexp.groupindex
  200. for m in matches:
  201. n = groupindex['quote'] - 1
  202. key = m[n]
  203. if key:
  204. quotes[key] = quotes.get(key, 0) + 1
  205. try:
  206. n = groupindex['delim'] - 1
  207. key = m[n]
  208. except KeyError:
  209. continue
  210. if key and (delimiters is None or key in delimiters):
  211. delims[key] = delims.get(key, 0) + 1
  212. try:
  213. n = groupindex['space'] - 1
  214. except KeyError:
  215. continue
  216. if m[n]:
  217. spaces += 1
  218. quotechar = max(quotes, key=quotes.get)
  219. if delims:
  220. delim = max(delims, key=delims.get)
  221. skipinitialspace = delims[delim] == spaces
  222. if delim == '\n': # most likely a file with a single column
  223. delim = ''
  224. else:
  225. # there is *no* delimiter, it's a single column of quoted data
  226. delim = ''
  227. skipinitialspace = 0
  228. # if we see an extra quote between delimiters, we've got a
  229. # double quoted format
  230. dq_regexp = re.compile(
  231. r"((%(delim)s)|^)\W*%(quote)s[^%(delim)s\n]*%(quote)s[^%(delim)s\n]*%(quote)s\W*((%(delim)s)|$)" % \
  232. {'delim':re.escape(delim), 'quote':quotechar}, re.MULTILINE)
  233. if dq_regexp.search(data):
  234. doublequote = True
  235. else:
  236. doublequote = False
  237. return (quotechar, doublequote, delim, skipinitialspace)
  238. def _guess_delimiter(self, data, delimiters):
  239. """
  240. The delimiter /should/ occur the same number of times on
  241. each row. However, due to malformed data, it may not. We don't want
  242. an all or nothing approach, so we allow for small variations in this
  243. number.
  244. 1) build a table of the frequency of each character on every line.
  245. 2) build a table of frequencies of this frequency (meta-frequency?),
  246. e.g. 'x occurred 5 times in 10 rows, 6 times in 1000 rows,
  247. 7 times in 2 rows'
  248. 3) use the mode of the meta-frequency to determine the /expected/
  249. frequency for that character
  250. 4) find out how often the character actually meets that goal
  251. 5) the character that best meets its goal is the delimiter
  252. For performance reasons, the data is evaluated in chunks, so it can
  253. try and evaluate the smallest portion of the data possible, evaluating
  254. additional chunks as necessary.
  255. """
  256. data = list(filter(None, data.split('\n')))
  257. ascii = [chr(c) for c in range(127)] # 7-bit ASCII
  258. # build frequency tables
  259. chunkLength = min(10, len(data))
  260. iteration = 0
  261. charFrequency = {}
  262. modes = {}
  263. delims = {}
  264. start, end = 0, chunkLength
  265. while start < len(data):
  266. iteration += 1
  267. for line in data[start:end]:
  268. for char in ascii:
  269. metaFrequency = charFrequency.get(char, {})
  270. # must count even if frequency is 0
  271. freq = line.count(char)
  272. # value is the mode
  273. metaFrequency[freq] = metaFrequency.get(freq, 0) + 1
  274. charFrequency[char] = metaFrequency
  275. for char in charFrequency.keys():
  276. items = list(charFrequency[char].items())
  277. if len(items) == 1 and items[0][0] == 0:
  278. continue
  279. # get the mode of the frequencies
  280. if len(items) > 1:
  281. modes[char] = max(items, key=lambda x: x[1])
  282. # adjust the mode - subtract the sum of all
  283. # other frequencies
  284. items.remove(modes[char])
  285. modes[char] = (modes[char][0], modes[char][1]
  286. - sum(item[1] for item in items))
  287. else:
  288. modes[char] = items[0]
  289. # build a list of possible delimiters
  290. modeList = modes.items()
  291. total = float(min(chunkLength * iteration, len(data)))
  292. # (rows of consistent data) / (number of rows) = 100%
  293. consistency = 1.0
  294. # minimum consistency threshold
  295. threshold = 0.9
  296. while len(delims) == 0 and consistency >= threshold:
  297. for k, v in modeList:
  298. if v[0] > 0 and v[1] > 0:
  299. if ((v[1]/total) >= consistency and
  300. (delimiters is None or k in delimiters)):
  301. delims[k] = v
  302. consistency -= 0.01
  303. if len(delims) == 1:
  304. delim = list(delims.keys())[0]
  305. skipinitialspace = (data[0].count(delim) ==
  306. data[0].count("%c " % delim))
  307. return (delim, skipinitialspace)
  308. # analyze another chunkLength lines
  309. start = end
  310. end += chunkLength
  311. if not delims:
  312. return ('', 0)
  313. # if there's more than one, fall back to a 'preferred' list
  314. if len(delims) > 1:
  315. for d in self.preferred:
  316. if d in delims.keys():
  317. skipinitialspace = (data[0].count(d) ==
  318. data[0].count("%c " % d))
  319. return (d, skipinitialspace)
  320. # nothing else indicates a preference, pick the character that
  321. # dominates(?)
  322. items = [(v,k) for (k,v) in delims.items()]
  323. items.sort()
  324. delim = items[-1][1]
  325. skipinitialspace = (data[0].count(delim) ==
  326. data[0].count("%c " % delim))
  327. return (delim, skipinitialspace)
  328. def has_header(self, sample):
  329. # Creates a dictionary of types of data in each column. If any
  330. # column is of a single type (say, integers), *except* for the first
  331. # row, then the first row is presumed to be labels. If the type
  332. # can't be determined, it is assumed to be a string in which case
  333. # the length of the string is the determining factor: if all of the
  334. # rows except for the first are the same length, it's a header.
  335. # Finally, a 'vote' is taken at the end for each column, adding or
  336. # subtracting from the likelihood of the first row being a header.
  337. rdr = reader(StringIO(sample), self.sniff(sample))
  338. header = next(rdr) # assume first row is header
  339. columns = len(header)
  340. columnTypes = {}
  341. for i in range(columns): columnTypes[i] = None
  342. checked = 0
  343. for row in rdr:
  344. # arbitrary number of rows to check, to keep it sane
  345. if checked > 20:
  346. break
  347. checked += 1
  348. if len(row) != columns:
  349. continue # skip rows that have irregular number of columns
  350. for col in list(columnTypes.keys()):
  351. thisType = complex
  352. try:
  353. thisType(row[col])
  354. except (ValueError, OverflowError):
  355. # fallback to length of string
  356. thisType = len(row[col])
  357. if thisType != columnTypes[col]:
  358. if columnTypes[col] is None: # add new column type
  359. columnTypes[col] = thisType
  360. else:
  361. # type is inconsistent, remove column from
  362. # consideration
  363. del columnTypes[col]
  364. # finally, compare results against first row and "vote"
  365. # on whether it's a header
  366. hasHeader = 0
  367. for col, colType in columnTypes.items():
  368. if isinstance(colType, int): # it's a length
  369. if len(header[col]) != colType:
  370. hasHeader += 1
  371. else:
  372. hasHeader -= 1
  373. else: # attempt typecast
  374. try:
  375. colType(header[col])
  376. except (ValueError, TypeError):
  377. hasHeader += 1
  378. else:
  379. hasHeader -= 1
  380. return hasHeader > 0