diff --git a/dataset/freeze/format/fcsv.py b/dataset/freeze/format/fcsv.py index e246904..0a69920 100644 --- a/dataset/freeze/format/fcsv.py +++ b/dataset/freeze/format/fcsv.py @@ -47,4 +47,5 @@ class CSVSerializer(Serializer): def close(self): for writer, fh in self.handles.values(): - fh.close() + if fh != self.fileobj: + fh.close() diff --git a/dataset/freeze/format/fjson.py b/dataset/freeze/format/fjson.py index 11a7714..b5783c3 100644 --- a/dataset/freeze/format/fjson.py +++ b/dataset/freeze/format/fjson.py @@ -55,4 +55,5 @@ class JSONSerializer(Serializer): self.export.get('callback'), data) fh.write(data) - fh.close() + if self.fileobj is None: + fh.close() diff --git a/test/test_freeze.py b/test/test_freeze.py index 09c8b98..82e6e07 100644 --- a/test/test_freeze.py +++ b/test/test_freeze.py @@ -62,6 +62,15 @@ class FreezeTestCase(unittest.TestCase): finally: fh.close() + def test_memory_streams(self): + import io + + for fmt in ('csv', 'json', 'tabson'): + with io.StringIO() as fd: + freeze(self.tbl.all(), format=fmt, fileobj=fd) + self.assertFalse(fd.closed, 'fileobj was closed for format %s' % fmt) + fd.getvalue() # should not throw + class SerializerTestCase(unittest.TestCase):