1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
# ~*~ coding: utf-8 ~*~
#
import unicodecsv
import codecs
from datetime import datetime
from six import BytesIO
from rest_framework.renderers import BaseRenderer
from rest_framework.utils import encoders, json
from ..utils import get_logger
logger = get_logger(__file__)
class JMSCSVRender(BaseRenderer):
media_type = 'text/csv'
format = 'csv'
@staticmethod
def _get_header(fields, template):
if template == 'import':
header = [
k for k, v in fields.items()
if not v.read_only and k != 'org_id'
]
elif template == 'update':
header = [k for k, v in fields.items() if not v.read_only]
else:
# template in ['export']
header = [k for k, v in fields.items() if not v.write_only]
return header
@staticmethod
def _gen_table(data, header, labels=None):
labels = labels or {}
yield [labels.get(k, k) for k in header]
for item in data:
row = [item.get(key) for key in header]
yield row
def set_response_disposition(self, serializer, context):
response = context.get('response')
if response and hasattr(serializer, 'Meta') and \
hasattr(serializer.Meta, "model"):
model_name = serializer.Meta.model.__name__.lower()
now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
filename = "{}_{}.csv".format(model_name, now)
disposition = 'attachment; filename="{}"'.format(filename)
response['Content-Disposition'] = disposition
def render(self, data, media_type=None, renderer_context=None):
renderer_context = renderer_context or {}
request = renderer_context['request']
template = request.query_params.get('template', 'export')
view = renderer_context['view']
if isinstance(data, dict) and data.get("count"):
data = data["results"]
if template == 'import':
data = [data[0]] if data else data
data = json.loads(json.dumps(data, cls=encoders.JSONEncoder))
try:
serializer = view.get_serializer()
self.set_response_disposition(serializer, renderer_context)
except Exception as e:
logger.debug(e, exc_info=True)
value = 'The resource not support export!'.encode('utf-8')
else:
fields = serializer.get_fields()
header = self._get_header(fields, template)
labels = {k: v.label for k, v in fields.items() if v.label}
table = self._gen_table(data, header, labels)
csv_buffer = BytesIO()
csv_buffer.write(codecs.BOM_UTF8)
csv_writer = unicodecsv.writer(csv_buffer, encoding='utf-8')
for row in table:
csv_writer.writerow(row)
value = csv_buffer.getvalue()
return value