Source code for horizon.utils.csvbase
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
from __future__ import division
from csv import DictWriter # noqa
from csv import writer # noqa
from django.http import HttpResponse # noqa
from django.http import StreamingHttpResponse # noqa
from django import template as django_template
import six
from six import StringIO
[docs]class CsvDataMixin(object):
"""CSV data Mixin - provides handling for CSV data.
.. attribute:: columns
A list of CSV column definitions. If omitted - no column titles
will be shown in the result file. Optional.
"""
def __init__(self):
self.out = StringIO()
super(CsvDataMixin, self).__init__()
if hasattr(self, "columns"):
columns = [self.encode(col) for col in self.columns]
self.writer = DictWriter(self.out, columns)
self.is_dict = True
else:
self.writer = writer(self.out)
self.is_dict = False
[docs] def write_csv_row(self, args):
if self.is_dict:
self.writer.writerow(dict(zip(
self.writer.fieldnames, [self.encode(col) for col in args])))
else:
self.writer.writerow([self.encode(col) for col in args])
[docs] def encode(self, value):
value = six.text_type(value)
if six.PY2:
# csv and StringIO cannot work with mixed encodings,
# so encode all with utf-8
value = value.encode('utf-8')
return value
[docs]class BaseCsvResponse(CsvDataMixin, HttpResponse):
"""Base CSV response class. Provides handling of CSV data."""
def __init__(self, request, template, context, content_type, **kwargs):
super(BaseCsvResponse, self).__init__()
self['Content-Disposition'] = 'attachment; filename="%s"' % (
kwargs.get("filename", "export.csv"),)
self['Content-Type'] = content_type
self.context = context
self.header = None
if template:
# Display some header info if provided as a template
header_template = django_template.loader.get_template(template)
context = django_template.RequestContext(request, self.context)
self.header = header_template.render(context)
if self.header:
self.out.write(self.encode(self.header))
self.write_csv_header()
for row in self.get_row_data():
self.write_csv_row(row)
self.out.flush()
self.content = self.out.getvalue()
self.out.close()
[docs] def get_row_data(self):
return []
[docs]class BaseCsvStreamingResponse(CsvDataMixin, StreamingHttpResponse):
"""Base CSV Streaming class. Provides streaming response for CSV data."""
def __init__(self, request, template, context, content_type, **kwargs):
super(BaseCsvStreamingResponse, self).__init__()
self['Content-Disposition'] = 'attachment; filename="%s"' % (
kwargs.get("filename", "export.csv"),)
self['Content-Type'] = content_type
self.context = context
self.header = None
if template:
# Display some header info if provided as a template
header_template = django_template.loader.get_template(template)
context = django_template.RequestContext(request, self.context)
self.header = header_template.render(context)
self._closable_objects.append(self.out)
self.streaming_content = self.get_content()
[docs] def buffer(self):
buf = self.out.getvalue()
self.out.truncate(0)
return buf
[docs] def get_content(self):
if self.header:
self.out.write(self.encode(self.header))
self.write_csv_header()
yield self.buffer()
for row in self.get_row_data():
self.write_csv_row(row)
yield self.buffer()
[docs] def get_row_data(self):
return []