Refactored CSV export logic

netbox/utilities/csv.py

@@ -0,0 +1,61 @@
+from __future__ import unicode_literals
+
+import datetime
+import six
+
+from django.http import HttpResponse
+
+
+def csv_format(data):
+    """
+    Encapsulate any data which contains a comma within double quotes.
+    """
+    csv = []
+    for value in data:
+
+        # Represent None or False with empty string
+        if value in [None, False]:
+            csv.append('')
+            continue
+
+        # Convert dates to ISO format
+        if isinstance(value, (datetime.date, datetime.datetime)):
+            value = value.isoformat()
+
+        # Force conversion to string first so we can check for any commas
+        if not isinstance(value, six.string_types):
+            value = '{}'.format(value)
+
+        # Double-quote the value if it contains a comma
+        if ',' in value:
+            csv.append('"{}"'.format(value))
+        else:
+            csv.append('{}'.format(value))
+
+    return ','.join(csv)
+
+
+def queryset_to_csv(queryset):
+    """
+    Export a queryset of objects as CSV, using the model's to_csv() method.
+    """
+    output = []
+
+    # Start with the column headers
+    headers = ','.join(queryset.model.csv_headers)
+    output.append(headers)
+
+    # Iterate through the queryset
+    for obj in queryset:
+        data = csv_format(obj.to_csv())
+        output.append(data)
+
+    # Build the HTTP response
+    response = HttpResponse(
+        '\n'.join(output),
+        content_type='text/csv'
+    )
+    filename = 'netbox_{}.csv'.format(queryset.model._meta.verbose_name_plural)
+    response['Content-Disposition'] = 'attachment; filename="{}"'.format(filename)
+
+    return response

netbox/utilities/utils.py

@@ -1,32 +1,5 @@
 from __future__ import unicode_literals
 
-import six
-
-
-def csv_format(data):
-    """
-    Encapsulate any data which contains a comma within double quotes.
-    """
-    csv = []
-    for value in data:
-
-        # Represent None or False with empty string
-        if value in [None, False]:
-            csv.append('')
-            continue
-
-        # Force conversion to string first so we can check for any commas
-        if not isinstance(value, six.string_types):
-            value = '{}'.format(value)
-
-        # Double-quote the value if it contains a comma
-        if ',' in value:
-            csv.append('"{}"'.format(value))
-        else:
-            csv.append('{}'.format(value))
-
-    return ','.join(csv)
-
 
 def foreground_color(bg_color):
     """

netbox/utilities/views.py

@@ -10,7 +10,6 @@ from django.core.exceptions import ValidationError
 from django.db import transaction, IntegrityError
 from django.db.models import ProtectedError
 from django.forms import CharField, Form, ModelMultipleChoiceField, MultipleHiddenInput, Textarea, TypedChoiceField
-from django.http import HttpResponse
 from django.shortcuts import get_object_or_404, redirect, render
 from django.template import TemplateSyntaxError
 from django.urls import reverse
@@ -21,6 +20,7 @@ from django.views.generic import View
 from django_tables2 import RequestConfig
 
 from extras.models import CustomField, CustomFieldValue, ExportTemplate, UserAction
+from utilities.csv import queryset_to_csv
 from utilities.forms import BootstrapMixin, CSVDataField
 from .error_handlers import handle_protectederror
 from .forms import ConfirmationForm
@@ -95,24 +95,15 @@ class ObjectListView(View):
             et = get_object_or_404(ExportTemplate, content_type=object_ct, name=request.GET.get('export'))
             queryset = CustomFieldQueryset(self.queryset, custom_fields) if custom_fields else self.queryset
             try:
-                response = et.to_response(context_dict={'queryset': queryset},
-                                          filename='netbox_{}'.format(model._meta.verbose_name_plural))
-                return response
+                return et.render_to_response(queryset)
             except TemplateSyntaxError:
-                messages.error(request, "There was an error rendering the selected export template ({})."
-                               .format(et.name))
-        # Fall back to built-in CSV export
+                messages.error(
+                    request,
+                    "There was an error rendering the selected export template ({}).".format(et.name)
+                )
+        # Fall back to built-in CSV export if no template was specified
         elif 'export' in request.GET and hasattr(model, 'to_csv'):
-            headers = getattr(model, 'csv_headers', None)
-            output = ','.join(headers) + '\n' if headers else ''
-            output += '\n'.join([obj.to_csv() for obj in self.queryset])
-            response = HttpResponse(
-                output,
-                content_type='text/csv'
-            )
-            response['Content-Disposition'] = 'attachment; filename="netbox_{}.csv"'\
-                .format(self.queryset.model._meta.verbose_name_plural)
-            return response
+            return queryset_to_csv(self.queryset)
 
         # Provide a hook to tweak the queryset based on the request immediately prior to rendering the object list
         self.queryset = self.alter_queryset(request)