1
|
|
|
# -*- coding: utf-8 -*- |
2
|
|
|
|
3
|
|
|
from benedict.serializers.abstract import AbstractSerializer |
4
|
|
|
from six import StringIO |
5
|
|
|
|
6
|
|
|
import csv |
7
|
|
|
|
8
|
|
|
|
9
|
|
|
class CSVSerializer(AbstractSerializer): |
10
|
|
|
|
11
|
|
|
def __init__(self): |
12
|
|
|
super(CSVSerializer, self).__init__() |
13
|
|
|
|
14
|
|
|
def decode(self, s, **kwargs): |
15
|
|
|
# kwargs.setdefault('delimiter', ',') |
16
|
|
|
if kwargs.pop('quote', False): |
17
|
|
|
kwargs.setdefault('quoting', csv.QUOTE_ALL) |
18
|
|
|
columns = kwargs.pop('columns', None) |
19
|
|
|
columns_row = kwargs.pop('columns_row', True) |
20
|
|
|
f = StringIO(s) |
21
|
|
|
r = csv.reader(f, **kwargs) |
22
|
|
|
ln = 0 |
23
|
|
|
data = [] |
24
|
|
|
for row in r: |
25
|
|
|
if ln == 0 and columns_row: |
26
|
|
|
if not columns: |
27
|
|
|
columns = row |
28
|
|
|
ln += 1 |
29
|
|
|
continue |
30
|
|
|
d = dict(zip(columns, row)) |
31
|
|
|
data.append(d) |
32
|
|
|
ln += 1 |
33
|
|
|
return data |
34
|
|
|
|
35
|
|
|
def encode(self, d, **kwargs): |
36
|
|
|
l = d |
37
|
|
|
# kwargs.setdefault('delimiter', ',') |
38
|
|
|
if kwargs.pop('quote', False): |
39
|
|
|
kwargs.setdefault('quoting', csv.QUOTE_ALL) |
40
|
|
|
kwargs.setdefault('lineterminator', '\n') |
41
|
|
|
columns = kwargs.pop('columns', None) |
42
|
|
|
columns_row = kwargs.pop('columns_row', True) |
43
|
|
|
if not columns and len(l) and isinstance(l[0], dict): |
44
|
|
|
keys = [str(key) for key in l[0].keys()] |
45
|
|
|
columns = list(sorted(keys)) |
46
|
|
|
f = StringIO() |
47
|
|
|
w = csv.writer(f, **kwargs) |
48
|
|
|
if columns_row and columns: |
49
|
|
|
w.writerow(columns) |
50
|
|
|
for item in l: |
51
|
|
|
if isinstance(item, dict): |
52
|
|
|
row = [item.get(key, '') for key in columns] |
53
|
|
|
elif isinstance(item, (list, tuple, set, )): |
54
|
|
|
row = item |
55
|
|
|
else: |
56
|
|
|
row = [item] |
57
|
|
|
w.writerow(row) |
58
|
|
|
data = f.getvalue() |
59
|
|
|
return data |
60
|
|
|
|