-
Notifications
You must be signed in to change notification settings - Fork 2k
/
writer.py
115 lines (94 loc) · 3.69 KB
/
writer.py
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# encoding: utf-8
from contextlib import contextmanager
from email.utils import encode_rfc2231
import json
import unicodecsv
UTF8_BOM = u'\uFEFF'.encode(u'utf-8')
@contextmanager
def csv_writer(response, columns, name=None, bom=False):
u'''Context manager for writing UTF-8 CSV data to response
:param response: file-like or response-like object for writing
data and headers (response-like objects only)
:param columns: list of column names
:param name: file name (for headers, response-like objects only)
:param bom: True to include a UTF-8 BOM at the start of the file
>>> with csv_writer(response, fields) as d:
>>> d.writerow(row1)
>>> d.writerow(row2)
'''
if hasattr(response, u'headers'):
response.headers['Content-Type'] = b'text/csv; charset=utf-8'
if name:
response.headers['Content-disposition'] = (
b'attachment; filename="{name}.csv"'.format(
name=encode_rfc2231(name)))
wr = unicodecsv.writer(response, encoding=u'utf-8')
if bom:
response.write(UTF8_BOM)
wr.writerow(columns)
yield wr
@contextmanager
def tsv_writer(response, columns, name=None, bom=False):
u'''Context manager for writing UTF-8 TSV data to response
:param response: file-like or response-like object for writing
data and headers (response-like objects only)
:param columns: list of column names
:param name: file name (for headers, response-like objects only)
:param bom: True to include a UTF-8 BOM at the start of the file
>>> with tsv_writer(response, fields) as d:
>>> d.writerow(row1)
>>> d.writerow(row2)
'''
if hasattr(response, u'headers'):
response.headers['Content-Type'] = (
b'text/tab-separated-values; charset=utf-8')
if name:
response.headers['Content-disposition'] = (
b'attachment; filename="{name}.tsv"'.format(
name=encode_rfc2231(name)))
wr = unicodecsv.writer(
response, encoding=u'utf-8', dialect=unicodecsv.excel_tab)
if bom:
response.write(UTF8_BOM)
wr.writerow(columns)
yield wr
@contextmanager
def json_writer(response, columns, name=None, bom=False):
u'''Context manager for writing UTF-8 JSON data to response
:param response: file-like or response-like object for writing
data and headers (response-like objects only)
:param columns: list of column names
:param name: file name (for headers, response-like objects only)
:param bom: True to include a UTF-8 BOM at the start of the file
>>> with json_writer(response, fields) as d:
>>> d.writerow(row1)
>>> d.writerow(row2)
'''
if hasattr(response, u'headers'):
response.headers['Content-Type'] = (
b'application/json; charset=utf-8')
if name:
response.headers['Content-disposition'] = (
b'attachment; filename="{name}.json"'.format(
name=encode_rfc2231(name)))
if bom:
response.write(UTF8_BOM)
response.write(b'{\n "data": [')
yield JSONWriter(response, columns)
response.write(b'\n]}\n')
class JSONWriter(object):
def __init__(self, response, columns):
self.response = response
self.columns = columns
self.first = True
def writerow(self, row):
if self.first:
self.first = False
self.response.write(b'\n ')
else:
self.response.write(b',\n ')
self.response.write(json.dumps(
{k: v for (k, v) in zip(self.columns, row)},
ensure_ascii=False,
separators=(u',', u':'),
sort_keys=True).encode(u'utf-8'))