Skip to content

Instantly share code, notes, and snippets.

@mjumbewu
Last active November 28, 2019 09:19
Show Gist options
  • Star 5 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save mjumbewu/4529220 to your computer and use it in GitHub Desktop.
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.
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 = '_'
#-*- 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