Merge pull request #62 from azavea/topic/support_aggregates

Topic/support aggregates

Author: Steve Lamb

djqscsv/__init__.py

@@ -1,2 +1,2 @@
 from djqscsv import (render_to_csv_response, write_csv,  # NOQA
-                             generate_filename, CSVException)  # NOQA
+                     generate_filename, CSVException)  # NOQA

djqscsv/_csql.py

@@ -6,12 +6,19 @@
 This module may later be officially supported.
 """
 
+def _identity(x):
+    return x
 
 def _transform(dataset, arg):
     if isinstance(arg, str):
-        return (dataset[0].index(arg), arg)
-    elif isinstance(arg, tuple):
-        return (dataset[0].index(arg[0]), arg[1])
+        field = arg
+        display_name = arg
+        transformer = _identity
+    else:
+        field, display_name, transformer = arg
+        if field is None:
+            field = dataset[0][0]
+    return (dataset[0].index(field), display_name, transformer)
 
 
 def SELECT(dataset, *args):
@@ -23,7 +30,7 @@ def SELECT(dataset, *args):
     results += [[header[1] for header in index_headers]]
 
     # add the rest of the rows
-    results += [[datarow[i] for i, h in index_headers]
+    results += [[trans(datarow[i]) for i, h, trans in index_headers]
                 for datarow in dataset[1:]]
     return results
 
@@ -34,5 +41,8 @@ def EXCLUDE(dataset, *args):
     return SELECT(dataset, *antiargs)
 
 
+def CONSTANT(value, display_name):
+    return (None, display_name, lambda x: value)
+    
 def AS(field, display_name):
-    return (field, display_name)
+    return (field, display_name, _identity)

djqscsv/djqscsv.py

@@ -90,6 +90,10 @@ def write_csv(queryset, file_obj, **kwargs):
     if extra_columns:
         field_names += extra_columns
 
+    aggregate_columns = list(values_qs.query.aggregate_select)
+    if aggregate_columns:
+        field_names += aggregate_columns
+
     if field_order:
         # go through the field_names and put the ones
         # that appear in the ordering list first

test_app/djqscsv_tests/context.py

@@ -6,4 +6,4 @@
 
 import djqscsv.djqscsv as djqscsv  # NOQA
 
-from djqscsv._csql import SELECT, EXCLUDE, AS
+from djqscsv._csql import SELECT, EXCLUDE, AS, CONSTANT

test_app/djqscsv_tests/tests/test_csv_creation.py

@@ -1,3 +1,4 @@
+from django.db.models import Count
 from django.test import TestCase
 from django.core.exceptions import ValidationError
 
@@ -8,7 +9,7 @@
 
 from djqscsv_tests.context import djqscsv
 
-from djqscsv_tests.context import SELECT, EXCLUDE, AS
+from djqscsv_tests.context import SELECT, EXCLUDE, AS, CONSTANT
 
 from djqscsv_tests.models import Person
 
@@ -213,6 +214,22 @@ def test_no_values_matches_models_file(self):
                                       use_verbose_names=False)
 
 
+class AggregateTests(CSVTestCase):
+
+    def setUp(self):
+        self.qs = create_people_and_get_queryset().annotate(num_hobbies=Count('hobby'))
+
+    def test_aggregate(self):
+        csv_with_aggregate = SELECT(self.FULL_PERSON_CSV,
+                                    'ID',
+                                    "Person's name",
+                                    'address',
+                                    "Info on Person",
+                                    'hobby_id',
+                                    CONSTANT('1', 'num_hobbies'))
+        self.assertQuerySetBecomesCsv(self.qs, csv_with_aggregate)
+
+
 class ExtraOrderingTests(CSVTestCase):
 
     def setUp(self):