Last active
November 28, 2019 09:19
-
-
Save mjumbewu/4529220 to your computer and use it in GitHub Desktop.
A Django REST Framework renderer which renders data from DRF serializers into CSV. The underlying functions will render any hierarchy of Python primitive containers to CSV.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import csv | |
from collections import defaultdict | |
from rest_framework.renderers import * | |
from StringIO import StringIO | |
class CSVRenderer(BaseRenderer): | |
""" | |
Renderer which serializes to CSV | |
""" | |
media_type = 'text/csv' | |
format = 'csv' | |
level_sep = '.' | |
def render(self, data, media_type=None, renderer_context=None): | |
""" | |
Renders serialized *data* into CSV. For a dictionary: | |
""" | |
if data is None: | |
return '' | |
table = self.tablize(data) | |
csv_buffer = StringIO() | |
csv_writer = csv.writer(csv_buffer) | |
for row in table: | |
# Assume that strings should be encoded as UTF-8 | |
csv_writer.writerow([ | |
elem.encode('utf-8') if isinstance(elem, basestring) else elem | |
for elem in row | |
]) | |
return csv_buffer.getvalue() | |
def tablize(self, data): | |
""" | |
Convert a list of data into a table. | |
""" | |
if data: | |
# First, flatten the data (i.e., convert it to a list of | |
# dictionaries that are each exactly one level deep). The key for | |
# each item designates the name of the column that the item will | |
# fall into. | |
data = self.flatten_data(data) | |
# Get the set of all unique headers, and sort them. | |
headers = set() | |
for item in data: | |
headers.update(item.keys()) | |
headers = sorted(headers) | |
# Create a row for each dictionary, filling in columns for which the | |
# item has no data with None values. | |
rows = [] | |
for item in data: | |
row = [] | |
for key in headers: | |
row.append(item.get(key, None)) | |
rows.append(row) | |
# Return your "table", with the headers as the first row. | |
return [headers] + rows | |
else: | |
return [] | |
def flatten_data(self, data): | |
""" | |
Convert the given data collection to a list of dictionaries that are | |
each exactly one level deep. The key for each value in the dictionaries | |
designates the name of the column that the value will fall into. | |
""" | |
flat_data = [] | |
for item in data: | |
flat_item = self.flatten_item(item) | |
flat_data.append(flat_item) | |
return flat_data | |
def flatten_item(self, item): | |
if isinstance(item, list): | |
flat_item = self.flatten_list(item) | |
elif isinstance(item, dict): | |
flat_item = self.flatten_dict(item) | |
else: | |
flat_item = {'': item} | |
return flat_item | |
def nest_flat_item(self, flat_item, prefix): | |
""" | |
Given a "flat item" (a dictionary exactly one level deep), nest all of | |
the column headers in a namespace designated by prefix. For example: | |
header... | with prefix... | becomes... | |
-----------|----------------|---------------- | |
'lat' | 'location' | 'location.lat' | |
'' | '0' | '0' | |
'votes.1' | 'user' | 'user.votes.1' | |
""" | |
nested_item = {} | |
for header, val in flat_item.iteritems(): | |
nested_header = self.level_sep.join([prefix, header]) if header else prefix | |
nested_item[nested_header] = val | |
return nested_item | |
def flatten_list(self, l): | |
flat_list = {} | |
for index, item in enumerate(l): | |
index = str(index) | |
flat_item = self.flatten_item(item) | |
nested_item = self.nest_flat_item(flat_item, index) | |
flat_list.update(nested_item) | |
return flat_list | |
def flatten_dict(self, d): | |
flat_dict = {} | |
for key, item in d.iteritems(): | |
key = str(key) | |
flat_item = self.flatten_item(item) | |
nested_item = self.nest_flat_item(flat_item, key) | |
flat_dict.update(nested_item) | |
return flat_dict | |
class CSVRendererWithUnderscores (CSVRenderer): | |
level_sep = '_' |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#-*- coding:utf-8 -*- | |
from django.test import TestCase | |
from ..renderers import CSVRenderer | |
class TestCSVRenderer (TestCase): | |
def test_tablize_a_list_with_no_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([]) | |
self.assertEqual(flat, []) | |
def test_tablize_a_list_with_atomic_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([1, 2, 'hello']) | |
self.assertEqual(flat, [['' ], | |
[1 ], | |
[2 ], | |
['hello']]) | |
def test_tablize_a_list_with_list_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([[1, 2, 3], | |
[4, 5], | |
[6, 7, [8, 9]]]) | |
self.assertEqual(flat, [['0' , '1' , '2' , '2.0' , '2.1'], | |
[1 , 2 , 3 , None , None ], | |
[4 , 5 , None , None , None ], | |
[6 , 7 , None , 8 , 9 ]]) | |
def test_tablize_a_list_with_dictionary_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([{'a': 1, 'b': 2}, | |
{'b': 3, 'c': {'x': 4, 'y': 5}}]) | |
self.assertEqual(flat, [['a' , 'b' , 'c.x' , 'c.y' ], | |
[1 , 2 , None , None ], | |
[None, 3 , 4 , 5 ]]) | |
def test_tablize_a_list_with_mixed_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([{'a': 1, 'b': 2}, | |
{'b': 3, 'c': [4, 5]}, | |
6]) | |
self.assertEqual(flat, [['' , 'a' , 'b' , 'c.0' , 'c.1'], | |
[None, 1 , 2 , None , None ], | |
[None, None, 3 , 4 , 5 ], | |
[6 , None, None, None , None ]]) | |
def test_tablize_a_list_with_unicode_elements(self): | |
renderer = CSVRenderer() | |
flat = renderer.tablize([{u'a': 1, u'b': u'hello\u2014goodbye'}]) | |
self.assertEqual(flat, [[u'a', u'b' ], | |
[1 , u'hello—goodbye']]) | |
def test_render_a_list_with_unicode_elements(self): | |
renderer = CSVRenderer() | |
dump = renderer.render([{u'a': 1, u'b': u'hello\u2014goodbye', u'c': 'http://example.com/'}]) | |
self.assertEqual(dump, (u'a,b,c\r\n1,hello—goodbye,http://example.com/\r\n').encode('utf-8')) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment