Merge pull request #337 from ZuluPro/stream

Added stream to CSV

3 files changed, 29 insertions, 3 deletions

diff --git a/.gitignore b/.gitignore
index 379075b..d6fdabd 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 8b536a7..9e8515a 100644
--- a/tablib/formats/_csv.py
+++ b/tablib/formats/_csv.py
@@ -13,8 +13,8 @@ extensions = ('csv',)
 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 e7b7233..6e3bc60 100755
--- a/test_tablib.py
+++ b/test_tablib.py
@@ -13,7 +13,7 @@ from uuid import uuid4
 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 @@ class TablibTestCase(unittest.TestCase):
 
         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."""