Merge pull request #106 from dimagi/sk/unicode-csv

encode headings in csv writer

commcare_export/writers.py

@@ -1,14 +1,14 @@
-import csv
 import datetime
+import io
 import logging
 import sys
 import zipfile
-from itertools import chain
 
 import alembic
+import csv342 as csv
 import six
 import sqlalchemy
-from six import StringIO, u
+from six import u
 
 logger = logging.getLogger(__name__)
 
@@ -82,17 +82,22 @@ def write_table(self, table):
         if self.archive is None:
             raise Exception('Attempt to write to a closed CsvWriter')
 
-        tempfile = StringIO()
+        def _encode_row(row):
+            return [
+                val.encode('utf-8') if isinstance(val, bytes) else val
+                for val in row
+            ]
+
+        tempfile = io.StringIO()
         writer = csv.writer(tempfile, dialect=csv.excel)
-        writer.writerow(table['headings'])
+        writer.writerow(_encode_row(table['headings']))
         for row in table['rows']:
-            writer.writerow([val.encode('utf-8') if isinstance(val, six.text_type) else val
-                             for val in row])
+            writer.writerow(_encode_row(row))
 
         # TODO: make this a polite zip and put everything in a subfolder with the same basename
         # as the zipfile
         self.archive.writestr('%s.csv' % self.zip_safe_name(table['name']),
-                              tempfile.getvalue())
+                              tempfile.getvalue().encode('utf-8'))
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.archive.close()

setup.py

@@ -79,7 +79,8 @@ def run_tests(self):
         'sqlalchemy',
         'pytz',
         'sqlalchemy-migrate',
-        'backoff'
+        'backoff',
+        'csv342'
     ],
     tests_require = ['pytest', 'psycopg2', 'mock'],
     cmdclass = {'test': PyTest},

tests/test_cli.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-import csv
+import csv342 as csv
 import os
 import unittest
 from argparse import Namespace

tests/test_writers.py

@@ -1,14 +1,17 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals, print_function, absolute_import, division, generators, nested_scopes
 
+import csv342 as csv
 import datetime
+import io
 import tempfile
+import zipfile
 
 import openpyxl
 import pytest
 import sqlalchemy
 
-from commcare_export.writers import SqlTableWriter, JValueTableWriter, Excel2007TableWriter
+from commcare_export.writers import SqlTableWriter, JValueTableWriter, Excel2007TableWriter, CsvTableWriter
 
 
 @pytest.fixture()
@@ -100,6 +103,28 @@ def _check_Excel2007TableWriter_output(self, filename):
                 ['4', '日本', '6'],
             ]
 
+    def test_CsvTableWriter(self):
+        with tempfile.NamedTemporaryFile() as file:
+            with CsvTableWriter(file=file) as writer:
+                writer.write_table({
+                    'name': 'foo',
+                    'headings': ['a', 'bjørn', 'c'],
+                    'rows': [
+                        [1, '2', 3],
+                        [4, '日本', 6],
+                    ]
+                })
+
+            with zipfile.ZipFile(file.name, 'r') as output_zip:
+                with output_zip.open('foo.csv') as csv_file:
+                    output = csv.reader(io.TextIOWrapper(csv_file, encoding='utf-8'))
+
+                    assert [row for row in output] == [
+                        ['a', 'bjørn', 'c'],
+                        ['1', '2', '3'],
+                        ['4', '日本', '6'],
+                    ]
+
 
 @pytest.mark.dbtest
 class TestSQLWriters(object):