diff --git a/.gitignore b/.gitignore index 379075b3..d6fdabdc 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,4 @@ junit-py27.xml # pyenv noise .python-version +tablib.egg-info/* diff --git a/tablib/formats/_csv.py b/tablib/formats/_csv.py index 8b536a7b..9e8515a4 100644 --- a/tablib/formats/_csv.py +++ b/tablib/formats/_csv.py @@ -13,8 +13,8 @@ DEFAULT_DELIMITER = unicode(',') -def export_set(dataset, **kwargs): - """Returns CSV representation of Dataset.""" +def export_stream_set(dataset, **kwargs): + """Returns CSV representation of Dataset as file-like.""" stream = StringIO() kwargs.setdefault('delimiter', DEFAULT_DELIMITER) @@ -24,6 +24,13 @@ def export_set(dataset, **kwargs): for row in dataset._package(dicts=False): _csv.writerow(row) + stream.seek(0) + return stream + + +def export_set(dataset, **kwargs): + """Returns CSV representation of Dataset.""" + stream = export_stream_set(dataset, **kwargs) return stream.getvalue() diff --git a/test_tablib.py b/test_tablib.py index e7b7233c..6e3bc608 100755 --- a/test_tablib.py +++ b/test_tablib.py @@ -13,7 +13,7 @@ import tablib from tablib.compat import markup, unicode, is_py3 from tablib.core import Row -from tablib.formats import csv as csv_format +from tablib.formats import _csv as csv_module class TablibTestCase(unittest.TestCase): @@ -262,6 +262,24 @@ def test_csv_export(self): self.assertEqual(csv, self.founders.csv) + def test_csv_stream_export(self): + """Verify exporting dataset object as CSV from file object.""" + + # Build up the csv string with headers first, followed by each row + csv = '' + for col in self.headers: + csv += col + ',' + + csv = csv.strip(',') + '\r\n' + + for founder in self.founders: + for col in founder: + csv += str(col) + ',' + csv = csv.strip(',') + '\r\n' + + csv_stream = csv_module.export_stream_set(self.founders) + self.assertEqual(csv, csv_stream.getvalue()) + def test_tsv_export(self): """Verify exporting dataset object as TSV."""