Fix python 2 support for StringIO.

Browse Source
This commit is contained in:
Friedrich Lindenberg 2016-01-18 10:56:35 +01:00
parent 5e9745b858
commit bd3a19ebc4
2 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
dataset/freeze/format/fcsv.py
View File

@ -25,7 +25,6 @@ class CSVSerializer(Serializer):
def write(self, path, result): def write(self, path, result):
keys = list(result.keys()) keys = list(result.keys())
if path not in self.handles: if path not in self.handles:
# handle fileobj that has been passed in: # handle fileobj that has been passed in:
if path is not None: if path is not None:
if PY3: # pragma: no cover if PY3: # pragma: no cover
@ -39,7 +38,7 @@ class CSVSerializer(Serializer):
if PY3: # pragma: no cover if PY3: # pragma: no cover
writer.writerow(keys) writer.writerow(keys)
else: else:
writer.writerow([k.encode('utf-8') for k in keys]) writer.writerow([value_to_str(k) for k in keys])
self.handles[path] = (writer, fh) self.handles[path] = (writer, fh)
writer, fh = self.handles[path] writer, fh = self.handles[path]
values = [value_to_str(result.get(k)) for k in keys] values = [value_to_str(result.get(k)) for k in keys]

9
test/test_freeze.py
View File

@ -63,13 +63,16 @@ class FreezeTestCase(unittest.TestCase):
fh.close() fh.close()
def test_memory_streams(self): def test_memory_streams(self):
import io if PY3:
from io import StringIO
else:
from io import BytesIO as StringIO
for fmt in ('csv', 'json', 'tabson'): for fmt in ('csv', 'json', 'tabson'):
with io.StringIO() as fd: with StringIO() as fd:
freeze(self.tbl.all(), format=fmt, fileobj=fd) freeze(self.tbl.all(), format=fmt, fileobj=fd)
self.assertFalse(fd.closed, 'fileobj was closed for format %s' % fmt) self.assertFalse(fd.closed, 'fileobj was closed for format %s' % fmt)
fd.getvalue() # should not throw fd.getvalue() # should not throw
class SerializerTestCase(unittest.TestCase): class SerializerTestCase(unittest.TestCase):