py3.py 2.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. # -*- coding: utf-8 -*-
  2. import csv
  3. from csv import *
  4. class _UnicodeWriteWrapper(object):
  5. """Simple write() wrapper that converts unicode to bytes."""
  6. def __init__(self, binary, encoding, errors):
  7. self.binary = binary
  8. self.encoding = encoding
  9. self.errors = errors
  10. def write(self, string):
  11. return self.binary.write(string.encode(self.encoding, self.errors))
  12. class UnicodeWriter(object):
  13. def __init__(self, f, dialect=csv.excel, encoding='utf-8', errors='strict',
  14. *args, **kwds):
  15. if f is None:
  16. raise TypeError
  17. f = _UnicodeWriteWrapper(f, encoding=encoding, errors=errors)
  18. self.writer = csv.writer(f, dialect, *args, **kwds)
  19. def writerow(self, row):
  20. return self.writer.writerow(row)
  21. def writerows(self, rows):
  22. return self.writer.writerows(rows)
  23. @property
  24. def dialect(self):
  25. return self.writer.dialect
  26. class UnicodeReader(object):
  27. def __init__(self, f, dialect=None, encoding='utf-8', errors='strict',
  28. **kwds):
  29. format_params = ['delimiter', 'doublequote', 'escapechar',
  30. 'lineterminator', 'quotechar', 'quoting',
  31. 'skipinitialspace']
  32. if dialect is None:
  33. if not any([kwd_name in format_params
  34. for kwd_name in kwds.keys()]):
  35. dialect = csv.excel
  36. f = (bs.decode(encoding, errors=errors) for bs in f)
  37. self.reader = csv.reader(f, dialect, **kwds)
  38. def __next__(self):
  39. return self.reader.__next__()
  40. def __iter__(self):
  41. return self
  42. @property
  43. def dialect(self):
  44. return self.reader.dialect
  45. @property
  46. def line_num(self):
  47. return self.reader.line_num
  48. writer = UnicodeWriter
  49. reader = UnicodeReader
  50. class DictWriter(csv.DictWriter):
  51. def __init__(self, csvfile, fieldnames, restval='',
  52. extrasaction='raise', dialect='excel', encoding='utf-8',
  53. errors='strict', *args, **kwds):
  54. super().__init__(csvfile, fieldnames, restval,
  55. extrasaction, dialect, *args, **kwds)
  56. self.writer = UnicodeWriter(csvfile, dialect, encoding=encoding,
  57. errors=errors, *args, **kwds)
  58. self.encoding_errors = errors
  59. def writeheader(self):
  60. header = dict(zip(self.fieldnames, self.fieldnames))
  61. self.writerow(header)
  62. class DictReader(csv.DictReader):
  63. def __init__(self, csvfile, fieldnames=None, restkey=None, restval=None,
  64. dialect='excel', encoding='utf-8', errors='strict', *args,
  65. **kwds):
  66. csv.DictReader.__init__(self, csvfile, fieldnames, restkey, restval,
  67. dialect, *args, **kwds)
  68. self.reader = UnicodeReader(csvfile, dialect, encoding=encoding,
  69. errors=errors, *args, **kwds)