csv.py 16 KB

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