csv.py 16 KB

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