ENH: can pass file handle or StringIO to Series.to_csv and DataFrame.to_csv, GH #765

Author: wesm

pandas/core/frame.py

@@ -909,15 +909,15 @@ def _helper_csvexcel(self, writer, na_rep=None, cols=None, header=True,
 
             writer.writerow(encoded_rows)
 
-    def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
-              index=True, index_label=None, mode='w', nanRep=None,
-              encoding=None):
+    def to_csv(self, path_or_buf, sep=",", na_rep='', cols=None,
+               header=True, index=True, index_label=None, mode='w',
+               nanRep=None, encoding=None):
         """
         Write DataFrame to a comma-separated values (csv) file
 
         Parameters
         ----------
-        path : string
+        path_or_buf : string or file handle / StringIO
             File path
         na_rep : string, default ''
             Missing data representation
@@ -944,11 +944,21 @@ def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
                           FutureWarning)
             na_rep = nanRep
 
-        with open(path, mode) as f:
+        if hasattr(path_or_buf, 'read'):
+            f = path_or_buf
+            close = False
+        else:
+            f = open(path_or_buf, mode)
+            close = True
+
+        try:
             csvout = csv.writer(f, lineterminator='\n', delimiter=sep)
             self._helper_csvexcel(csvout, na_rep=na_rep, cols=cols,
                                   header=header, index=index,
                                   index_label=index_label, encoding=encoding)
+        finally:
+            if close:
+                f.close()
 
     def to_excel(self, excel_writer, sheet_name = 'sheet1', na_rep='', cols=None, header=True,
                  index=True, index_label=None):

pandas/core/series.py

@@ -2056,7 +2056,7 @@ def to_csv(self, path, index=True, sep=",", na_rep='', header=False,
         from pandas.core.frame import DataFrame
         df = DataFrame(self)
         df.to_csv(path, index=index, sep=sep, na_rep=na_rep, header=header,
-                  index_label=index_label,mode=mode, nanRep=nanRep,
+                  index_label=index_label, mode=mode, nanRep=nanRep,
                   encoding=encoding)
 
     def dropna(self):

pandas/tests/test_frame.py

@@ -2254,32 +2254,37 @@ def test_to_csv_withcommas(self):
         os.remove(path)
 
     def test_to_csv_bug(self):
-        from pandas import read_csv
         path = '__tmp__.csv'
         f1 = StringIO('a,1.0\nb,2.0')
         df = DataFrame.from_csv(f1,header=None)
         newdf = DataFrame({'t': df[df.columns[0]]})
         newdf.to_csv(path)
 
-        recons = read_csv(path, index_col=0)
+        recons = pan.read_csv(path, index_col=0)
         assert_frame_equal(recons, newdf)
 
         os.remove(path)
 
     def test_to_csv_unicode(self):
-        from pandas import read_csv
         path = '__tmp__.csv'
         df = DataFrame({u'c/\u03c3':[1,2,3]})
         df.to_csv(path, encoding='UTF-8')
-        df2 = read_csv(path, index_col=0, encoding='UTF-8')
+        df2 = pan.read_csv(path, index_col=0, encoding='UTF-8')
         assert_frame_equal(df, df2)
 
         df.to_csv(path, encoding='UTF-8', index=False)
-        df2 = read_csv(path, index_col=None, encoding='UTF-8')
+        df2 = pan.read_csv(path, index_col=None, encoding='UTF-8')
         assert_frame_equal(df, df2)
 
         os.remove(path)
 
+    def test_to_csv_stringio(self):
+        buf = StringIO()
+        self.frame.to_csv(buf)
+        buf.seek(0)
+        recons = pan.read_csv(buf, index_col=0)
+        assert_frame_equal(recons, self.frame)
+
     def test_to_excel_from_excel(self):
         try:
             import xlwt

pandas/tests/test_series.py

@@ -1,5 +1,6 @@
 # pylint: disable-msg=E1101,W0612
 
+from cStringIO import StringIO
 from datetime import datetime, timedelta
 import os
 import operator
@@ -1515,6 +1516,13 @@ def test_to_csv(self):
 
         os.remove('_foo')
 
+    def test_to_csv_stringio(self):
+        buf = StringIO()
+        self.ts.to_csv(buf, index=False)
+        buf.seek(0)
+        arr = np.loadtxt(buf)
+        assert_almost_equal(arr, self.ts.values)
+
     def test_to_dict(self):
         self.assert_(np.array_equal(Series(self.ts.to_dict()), self.ts))