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