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