actions: change export_csv to also exports object and string data in columns

django_journal/actions.py

@@ -1,23 +1,41 @@
 import csv
 
+
 from django.utils.translation import ugettext_lazy as _
 from django.http import HttpResponse
 
+
+from . import models
+
+def export_as_csv_generator(queryset):
+    header = ['time', 'tag', 'message']
+    tags = set(models.Tag.objects.filter(objectdata__journal__in=queryset).values_list('name', flat=True))
+    tags |= set(models.Tag.objects.filter(stringdata__journal__in=queryset).values_list('name', flat=True))
+    extra_headers = map(lambda s: s.encode('utf-8'), sorted(tags))
+    yield header+extra_headers
+    for journal in queryset:
+        row = {
+                'time': journal.time.isoformat(' '),
+                'tag': journal.tag.name.encode('utf-8'),
+                'message': unicode(journal).encode('utf-8'),
+              }
+        for stringdata in journal.stringdata_set.all():
+            row[stringdata.tag.name.encode('utf-8')] = stringdata.content.encode('utf-8')
+        for objectdata in journal.objectdata_set.all():
+            row[objectdata.tag.name.encode('utf-8')] = str(objectdata.object_id)
+        yield row
+
 def export_as_csv(modeladmin, request, queryset):
     """
     CSV export for journal
     """
     response = HttpResponse(mimetype='text/csv')
     response['Content-Disposition'] = 'attachment; filename=journal.csv'
-    writer = csv.writer(response)
-
-    writer.writerow(['time', 'tag', 'message'])
-    for journal in queryset:
-        row = [
-                journal.time.isoformat(' '),
-                journal.tag.name.encode('utf-8'),
-                unicode(journal).encode('utf-8'),
-              ]
+    l = export_as_csv_generator(queryset)
+    header = l.next()
+    writer = csv.DictWriter(response, header)
+    writer.writerow(dict(zip(header, header)))
+    for row in l:
         writer.writerow(row)
     return response
 export_as_csv.short_description = _(u"Export CSV file")

django_journal/tests.py

@@ -2,8 +2,11 @@ from django.test import TestCase
 from django.contrib.auth.models import User, Group
 from django.db import transaction
 
+
 from . import record
-import models
+from . import actions
+from . import models
+
 
 class JournalTestCase(TestCase):
     def setUp(self):
@@ -39,3 +42,10 @@ class JournalTestCase(TestCase):
         for i, event in zip(range(20), models.Journal.objects.for_tag('logout').order_by('id')):
             self.assertEqual(unicode(event), 'user{0} logged out'.format(i))
 
+    def test_export_as_csv(self):
+        qs = models.Journal.objects.all()
+        l = list(actions.export_as_csv_generator(qs))
+        assert l[0] == ['time', 'tag', 'message', 'group', 'user', 'user1', 'user2']
+        l = list(actions.export_as_csv_generator(qs[:5]))
+        assert l[0] == ['time', 'tag', 'message', 'user']
+