finally support encoding for csv output

Author: chfw

pyexcel_io/fileformat/_csv.py

@@ -53,6 +53,36 @@ def __next__(self):
         return next(self.reader).encode('utf-8')
 
 
+class UnicodeWriter:
+    """
+    A CSV writer which will write rows to CSV file "f",
+    which is encoded in the given encoding.
+    """
+
+    def __init__(self, f, encoding="utf-8", **kwds):
+        # Redirect output to a queue
+        self.queue = StringIO()
+        self.writer = csv.writer(self.queue, **kwds)
+        self.stream = f
+        self.encoder = codecs.getincrementalencoder(encoding)()
+
+    def writerow(self, row):
+        self.writer.writerow([text_type(s).encode("utf-8") for s in row])
+        # Fetch UTF-8 output from the queue ...
+        data = self.queue.getvalue()
+        data = data.decode("utf-8")
+        # ... and reencode it into the target encoding
+        data = self.encoder.encode(data)
+        # write to the target stream
+        self.stream.write(data)
+        # empty queue
+        self.queue.truncate(0)
+
+    def writerows(self, rows):
+        for row in rows:
+            self.writerow(row)
+
+
 class CSVSheetReader(SheetReader):
     def __init__(self, sheet, encoding="utf-8",
                  auto_detect_float=True, auto_detect_int=True,
@@ -149,12 +179,7 @@ def write_row(self, array):
         """
         write a row into the file
         """
-        if PY2:
-            self.writer.writerow(
-                [text_type(s if s is not None else '').encode(self.encoding)
-                 for s in array])
-        else:
-            self.writer.writerow(array)
+        self.writer.writerow(array)
 
 
 class CSVFileWriter(CSVSheetWriter):
@@ -175,9 +200,12 @@ def set_sheet_name(self, name):
             file_name = self.native_book
         if PY2:
             self.f = open(file_name, "wb")
+            self.writer = UnicodeWriter(self.f, encoding=self.encoding,
+                                        **self.keywords)
         else:
-            self.f = open(file_name, "w", newline="")
-        self.writer = csv.writer(self.f, **self.keywords)
+            self.f = open(file_name, "w", newline="",
+                          encoding=self.encoding)
+            self.writer = csv.writer(self.f, **self.keywords)
 
 
 class CSVMemoryWriter(CSVSheetWriter):
@@ -194,8 +222,13 @@ def __init__(self, filename, name,
                 self.line_terminator))
 
     def set_sheet_name(self, name):
-        self.f = self.native_book
-        self.writer = csv.writer(self.f, **self.keywords)
+        if PY2:
+            self.f = self.native_book
+            self.writer = UnicodeWriter(self.f, encoding=self.encoding,
+                                        **self.keywords)
+        else:
+            self.f = self.native_book
+            self.writer = csv.writer(self.f, **self.keywords)
 
     def close(self):
         if self.single_sheet_in_book:

tests/test_csv_book.py

@@ -14,7 +14,7 @@
     CSVFileWriter,
     CSVMemoryWriter
 )
-from pyexcel_io._compact import BytesIO, PY2
+from pyexcel_io._compact import BytesIO, PY2, StringIO
 
 
 class TestReaders(TestCase):
@@ -154,15 +154,18 @@ def test_utf16_encoding():
     test_file = "test-utf16-encoding.csv"
     writer = CSVFileWriter(
         test_file, None,
-        encoding="utf-16")
+        encoding="utf-16", lineterminator="\n")
     writer.write_array(content)
     writer.close()
-    with open(test_file, "r") as f:
-        actual = f.read()
-        print(actual)
+    with open(test_file, "rb") as f:
+        actual = f.read().decode('utf-16')
+        if PY2:
+            actual = actual.encode('utf-8')
+        eq_(actual, 'Äkkilähdöt,Matkakirjoituksia,Matkatoimistot\n')
+    os.unlink(test_file)
 
 
-def test_utf16_memory_encoding():
+def test_utf16_memory_decoding():
     test_content = u'Äkkilähdöt,Matkakirjoituksia,Matkatoimistot'
     test_content = BytesIO(test_content.encode('utf-16'))
     reader = CSVinMemoryReader(
@@ -173,3 +176,16 @@ def test_utf16_memory_encoding():
         content[0] = [s.encode('utf-8') for s in content[0]]
     expected = [['Äkkilähdöt', 'Matkakirjoituksia', 'Matkatoimistot']]
     eq_(content, expected)
+
+
+def test_utf16_memory_encoding():
+    content = [[u'Äkkilähdöt', u'Matkakirjoituksia', u'Matkatoimistot']]
+    io = StringIO()
+    writer = CSVMemoryWriter(
+        io, None, lineterminator="\n", single_sheet_in_book=True,
+        encoding="utf-16")
+    writer.write_array(content)
+    actual = io.getvalue()
+    if PY2:
+        actual = actual.decode('utf-16')
+    eq_(actual, u'Äkkilähdöt,Matkakirjoituksia,Matkatoimistot\n')