Introduce common base class for CSV exports

src/pretix/base/exporter.py

@@ -1,5 +1,11 @@
+import io
+from collections import OrderedDict
 from typing import Tuple
 
+from defusedcsv import csv
+from django import forms
+from django.utils.translation import ugettext_lazy as _
+
 
 class BaseExporter:
     """
@@ -69,3 +75,46 @@ class BaseExporter:
         tasks.
         """
         raise NotImplementedError()  # NOQA
+
+
+class ListExporter(BaseExporter):
+
+    @property
+    def export_form_fields(self) -> dict:
+        ff = OrderedDict(
+            [
+                ('_format',
+                 forms.ChoiceField(
+                     label=_('Export format'),
+                     choices=(
+                         ('default', _('CSV (with commas)')),
+                         ('excel', _('CSV (Excel-style)')),
+                         ('semicolon', _('CSV (with semicolons)')),
+                     ),
+                 )),
+            ]
+        )
+        ff.update(self.additional_form_fields)
+        return ff
+
+    @property
+    def additional_form_fields(self) -> dict:
+        return {}
+
+    def iterate_list(self, form_data):
+        raise NotImplementedError()  # noqa
+
+    def get_filename(self):
+        return 'export.csv'
+
+    def render(self, form_data: dict) -> Tuple[str, str, str]:
+        output = io.StringIO()
+        if form_data.get('_format') == 'default':
+            writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC, delimiter=",")
+        elif form_data.get('_format') == 'excel':
+            writer = csv.writer(output, dialect='excel')
+        elif form_data.get('_format') == 'semicolon':
+            writer = csv.writer(output, dialect='excel', delimiter=";")
+        for line in self.iterate_list(form_data):
+            writer.writerow(line)
+        return self.get_filename(), 'text/csv', output.getvalue().encode("utf-8")

src/pretix/base/exporters/orderlist.py

@@ -1,9 +1,7 @@
-import io
 from collections import OrderedDict
 from decimal import Decimal
 
 import pytz
-from defusedcsv import csv
 from django import forms
 from django.db.models import DateTimeField, Max, OuterRef, Subquery, Sum
 from django.dispatch import receiver
@@ -14,16 +12,16 @@ from pretix.base.models import InvoiceAddress, Order, OrderPosition
 from pretix.base.models.orders import OrderFee, OrderPayment, OrderRefund
 from pretix.base.settings import PERSON_NAME_SCHEMES
 
-from ..exporter import BaseExporter
+from ..exporter import ListExporter
 from ..signals import register_data_exporters
 
 
-class OrderListExporter(BaseExporter):
+class OrderListExporter(ListExporter):
-    identifier = 'orderlistcsv'
+    identifier = 'orderlist'
-    verbose_name = ugettext_lazy('List of orders (CSV)')
+    verbose_name = ugettext_lazy('List of orders')
 
     @property
-    def export_form_fields(self):
+    def additional_form_fields(self):
         return OrderedDict(
             [
                 ('paid_only',
@@ -51,10 +49,8 @@ class OrderListExporter(BaseExporter):
         tax_rates = sorted(tax_rates)
         return tax_rates
 
-    def render(self, form_data: dict):
+    def iterate_list(self, form_data: dict):
-        output = io.StringIO()
         tz = pytz.timezone(self.event.settings.timezone)
-        writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC, delimiter=",")
 
         p_date = OrderPayment.objects.filter(
             order=OuterRef('pk'),
@@ -95,7 +91,7 @@ class OrderListExporter(BaseExporter):
 
         headers.append(_('Invoice numbers'))
 
-        writer.writerow(headers)
+        yield headers
 
         full_fee_sum_cache = {
             o['order__id']: o['grosssum'] for o in
@@ -162,17 +158,18 @@ class OrderListExporter(BaseExporter):
                 ]
 
             row.append(', '.join([i.number for i in order.invoices.all()]))
-            writer.writerow(row)
+            yield row
 
-        return '{}_orders.csv'.format(self.event.slug), 'text/csv', output.getvalue().encode("utf-8")
+    def get_filename(self):
+        return '{}_orders.csv'.format(self.event.slug)
 
 
-class PaymentListExporter(BaseExporter):
+class PaymentListExporter(ListExporter):
-    identifier = 'paymentlistcsv'
+    identifier = 'paymentlist'
-    verbose_name = ugettext_lazy('List of payments and refunds (CSV)')
+    verbose_name = ugettext_lazy('List of payments and refunds')
 
     @property
-    def export_form_fields(self):
+    def additional_form_fields(self):
         return OrderedDict(
             [
                 ('successful_only',
@@ -184,10 +181,8 @@ class PaymentListExporter(BaseExporter):
             ]
         )
 
-    def render(self, form_data: dict):
+    def iterate_list(self, form_data):
-        output = io.StringIO()
         tz = pytz.timezone(self.event.settings.timezone)
-        writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC, delimiter=",")
 
         provider_names = {
             k: v.verbose_name
@@ -215,7 +210,7 @@ class PaymentListExporter(BaseExporter):
             _('Order'), _('Payment ID'), _('Creation date'), _('Completion date'), _('Status'),
             _('Amount'), _('Payment method')
         ]
-        writer.writerow(headers)
+        yield headers
 
         for obj in objs:
             if isinstance(obj, OrderPayment) and obj.payment_date:
@@ -233,24 +228,22 @@ class PaymentListExporter(BaseExporter):
                 localize(obj.amount * (-1 if isinstance(obj, OrderRefund) else 1)),
                 provider_names.get(obj.provider, obj.provider)
             ]
-            writer.writerow(row)
+            yield row
 
-        return '{}_payments.csv'.format(self.event.slug), 'text/csv', output.getvalue().encode("utf-8")
+    def get_filename(self):
+        return '{}_payments.csv'.format(self.event.slug)
 
 
-class QuotaListExporter(BaseExporter):
+class QuotaListExporter(ListExporter):
-    identifier = 'quotalistcsv'
+    identifier = 'quotalist'
-    verbose_name = ugettext_lazy('Quota availabilities (CSV)')
+    verbose_name = ugettext_lazy('Quota availabilities')
-
-    def render(self, form_data: dict):
-        output = io.StringIO()
-        writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC, delimiter=",")
 
+    def iterate_list(self, form_data):
         headers = [
             _('Quota name'), _('Total quota'), _('Paid orders'), _('Pending orders'), _('Blocking vouchers'),
             _('Current user\'s carts'), _('Waiting list'), _('Current availability')
         ]
-        writer.writerow(headers)
+        yield headers
 
         for quota in self.event.quotas.all():
             avail = quota.availability()
@@ -264,9 +257,10 @@ class QuotaListExporter(BaseExporter):
                 quota.count_waiting_list_pending(),
                 _('Infinite') if avail[1] is None else avail[1]
             ]
-            writer.writerow(row)
+            yield row
 
-        return '{}_quotas.csv'.format(self.event.slug), 'text/csv', output.getvalue().encode("utf-8")
+    def get_filename(self):
+        return '{}_quotas.csv'.format(self.event.slug)
 
 
 @receiver(register_data_exporters, dispatch_uid="exporter_orderlist")

src/pretix/plugins/checkinlists/exporters.py

@@ -1,8 +1,6 @@
-import io
 from collections import OrderedDict
 
 import dateutil.parser
-from defusedcsv import csv
 from django import forms
 from django.conf import settings
 from django.db.models import Max, OuterRef, Subquery
@@ -16,7 +14,7 @@ from pytz import UTC
 from reportlab.lib.units import mm
 from reportlab.platypus import Flowable, Paragraph, Spacer, Table, TableStyle
 
-from pretix.base.exporter import BaseExporter
+from pretix.base.exporter import BaseExporter, ListExporter
 from pretix.base.models import (
     Checkin, InvoiceAddress, Order, OrderPosition, Question,
 )
@@ -26,9 +24,9 @@ from pretix.control.forms.widgets import Select2
 from pretix.plugins.reports.exporters import ReportlabExportMixin
 
 
-class BaseCheckinList(BaseExporter):
+class CheckInListMixin(BaseExporter):
     @property
-    def export_form_fields(self):
+    def _fields(self):
         name_scheme = PERSON_NAME_SCHEMES[self.event.settings.name_scheme]
         d = OrderedDict(
             [
@@ -155,14 +153,14 @@ class TableTextRotate(Flowable):
         canvas.drawString(0, -1, self.text)
 
 
-class PDFCheckinList(ReportlabExportMixin, BaseCheckinList):
+class PDFCheckinList(ReportlabExportMixin, CheckInListMixin, BaseExporter):
     name = "overview"
     identifier = 'checkinlistpdf'
     verbose_name = ugettext_lazy('Check-in list (PDF)')
 
     @property
     def export_form_fields(self):
-        f = super().export_form_fields
+        f = self._fields
         del f['secrets']
         return f
 
@@ -276,33 +274,16 @@ class PDFCheckinList(ReportlabExportMixin, BaseCheckinList):
         return story
 
 
-class CSVCheckinList(BaseCheckinList):
+class CSVCheckinList(CheckInListMixin, ListExporter):
     name = "overview"
-    identifier = 'checkinlistcsv'
+    identifier = 'checkinlist'
-    verbose_name = ugettext_lazy('Check-in list (CSV)')
+    verbose_name = ugettext_lazy('Check-in list')
 
     @property
-    def export_form_fields(self):
+    def additional_form_fields(self):
-        d = super().export_form_fields
+        return self._fields
-        d['dialect'] = forms.ChoiceField(
-            label=_('CSV dialect'),
-            choices=(
-                ('default', 'Default'),
-                ('excel', 'Excel'),
-                ('semicolon', 'Semicolon'),
-            )
-        )
-        return d
-
-    def render(self, form_data: dict):
-        output = io.StringIO()
-        if form_data.get('dialect', '-') in csv.list_dialects():
-            writer = csv.writer(output, dialect=form_data.get('dialect'))
-        elif form_data.get('dialect', '-') == "semicolon":
-            writer = csv.writer(output, dialect='excel', delimiter=';')
-        else:
-            writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC, delimiter=",")
 
+    def iterate_list(self, form_data):
         cl = self.event.checkin_lists.get(pk=form_data['list'])
 
         questions = list(Question.objects.filter(event=self.event, id__in=form_data['questions']))
@@ -339,7 +320,7 @@ class CSVCheckinList(BaseCheckinList):
 
         headers.append(_('Company'))
         headers.append(_('Voucher code'))
-        writer.writerow(headers)
+        yield headers
 
         for op in qs:
             try:
@@ -391,6 +372,7 @@ class CSVCheckinList(BaseCheckinList):
 
             row.append(ia.company)
             row.append(op.voucher.code if op.voucher else "")
-            writer.writerow(row)
+            yield row
 
-        return '{}_checkin.csv'.format(self.event.slug), 'text/csv', output.getvalue().encode("utf-8")
+    def get_filename(self):
+        return '{}_checkin.csv'.format(self.event.slug)