BUG: made encoding optional on csv read/write, addresses #717

Author: adamklein

pandas/core/common.py

@@ -831,11 +831,11 @@ def console_encode(value):
     except (AttributeError, TypeError):
         return value.encode('ascii', 'replace')
 
-def csv_encode(value):
+def csv_encode(value, encoding='UTF-8'):
     if py3compat.PY3 or not isinstance(value, unicode):
         return value
 
-    return value.encode('UTF-8', 'replace')
+    return value.encode(encoding, 'replace')
 
 class UTF8Recoder:
     """

pandas/core/frame.py

@@ -834,7 +834,8 @@ def to_panel(self):
     to_wide = deprecate('to_wide', to_panel)
 
     def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
-              index=True, index_label=None, mode='w', nanRep=None):
+              index=True, index_label=None, mode='w', nanRep=None, 
+              encoding=None):
         """
         Write DataFrame to a comma-separated values (csv) file
 
@@ -891,11 +892,23 @@ def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
                     # given a string for a DF with Index
                     index_label = [index_label]
 
-                encoded_labels = [csv_encode(val) for val in index_label]
-                encoded_cols = [csv_encode(val) for val in cols]
+                if encoding is not None:
+                    encoded_labels = [csv_encode(val, encoding=encoding)
+                                      for val in index_label]
+                    encoded_cols = [csv_encode(val, encoding=encoding)
+                                    for val in cols]
+                else:
+                    encoded_labels = list(index_label)
+                    encoded_cols = list(cols)
+
                 csvout.writerow(encoded_labels + encoded_cols)
             else:
-                encoded_cols = [csv_encode(val) for val in cols]
+                if encoding is not None:
+                    encoded_cols = [csv_encode(val, encoding=encoding)
+                                    for val in cols]
+                else:
+                    encoded_cols = list(cols)
+
                 csvout.writerow(encoded_cols)
 
         nlevels = getattr(self.index, 'nlevels', 1)
@@ -913,7 +926,12 @@ def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
 
                 row_fields.append(val)
 
-            encoded_rows = [csv_encode(val) for val in row_fields]
+            if encoding is not None:
+                encoded_rows = [csv_encode(val, encoding=encoding) 
+                                for val in row_fields]
+            else:
+                encoded_rows = list(row_fields)
+
             csvout.writerow(encoded_rows)
 
         f.close()

pandas/io/parsers.py

@@ -86,7 +86,8 @@
 def read_csv(filepath_or_buffer, sep=',', header=0, index_col=None, names=None,
              skiprows=None, na_values=None, parse_dates=False,
              date_parser=None, nrows=None, iterator=False, chunksize=None,
-             skip_footer=0, converters=None, verbose=False, delimiter=None):
+             skip_footer=0, converters=None, verbose=False, delimiter=None,
+             encoding=None):
     if hasattr(filepath_or_buffer, 'read'):
         f = filepath_or_buffer
     else:
@@ -111,7 +112,8 @@ def read_csv(filepath_or_buffer, sep=',', header=0, index_col=None, names=None,
                         chunksize=chunksize,
                         skip_footer=skip_footer,
                         converters=converters,
-                        verbose=verbose)
+                        verbose=verbose,
+                        encoding=encoding)
 
     if nrows is not None:
         return parser.get_chunk(nrows)
@@ -124,14 +126,15 @@ def read_csv(filepath_or_buffer, sep=',', header=0, index_col=None, names=None,
 def read_table(filepath_or_buffer, sep='\t', header=0, index_col=None,
                names=None, skiprows=None, na_values=None, parse_dates=False,
                date_parser=None, nrows=None, iterator=False, chunksize=None,
-               skip_footer=0, converters=None, verbose=False, delimiter=None):
+               skip_footer=0, converters=None, verbose=False, delimiter=None,
+               encoding=None):
     return read_csv(filepath_or_buffer, sep=sep, header=header,
                     skiprows=skiprows, index_col=index_col,
                     na_values=na_values, date_parser=date_parser,
                     names=names, parse_dates=parse_dates,
                     nrows=nrows, iterator=iterator, chunksize=chunksize,
                     skip_footer=skip_footer, converters=converters,
-                    verbose=verbose, delimiter=delimiter)
+                    verbose=verbose, delimiter=delimiter, encoding=None)
 
 def read_clipboard(**kwargs):  # pragma: no cover
     """
@@ -194,7 +197,8 @@ class TextParser(object):
     def __init__(self, f, delimiter=None, names=None, header=0,
                  index_col=None, na_values=None, parse_dates=False,
                  date_parser=None, chunksize=None, skiprows=None,
-                 skip_footer=0, converters=None, verbose=False):
+                 skip_footer=0, converters=None, verbose=False,
+                 encoding=None):
         """
         Workhorse function for processing nested list into DataFrame
 
@@ -210,6 +214,8 @@ def __init__(self, f, delimiter=None, names=None, header=0,
         self.date_parser = date_parser
         self.chunksize = chunksize
         self.passed_names = names is not None
+        self.encoding = encoding
+        
 
         if com.is_integer(skiprows):
             skiprows = range(skiprows)
@@ -261,9 +267,20 @@ def _make_reader(self, f):
                 self.pos += 1
                 sniffed = csv.Sniffer().sniff(line)
                 dia.delimiter = sniffed.delimiter
-                self.buf.extend(list(com.UnicodeReader(StringIO(line),
-                                dialect=dia)))
-            reader = com.UnicodeReader(f, dialect=dia)
+                if self.encoding is not None:
+                    self.buf.extend(list(
+                        com.UnicodeReader(StringIO(line),
+                                          dialect=dia, 
+                                          encoding=self.encoding)))
+                else:
+                    self.buf.extend(list(csv.reader(StringIO(line), 
+                                                    dialect=dia)))
+
+            if self.encoding is not None:
+                reader = com.UnicodeReader(f, dialect=dia, 
+                                           encoding=self.encoding)
+            else:
+                reader = csv.reader(f, dialect=dia)
         else:
             reader = (re.split(sep, line.strip()) for line in f)
 

pandas/tests/test_frame.py

@@ -24,7 +24,6 @@
                                  assert_frame_equal)
 
 import pandas.util.testing as tm
-from pandas.util import py3compat
 import pandas._tseries as lib
 
 #-------------------------------------------------------------------------------
@@ -2486,8 +2485,8 @@ 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)
-        df2 = read_csv(path, index_col=0)
+        df.to_csv(path, encoding='UTF-8')
+        df2 = read_csv(path, index_col=0, encoding='UTF-8')
         assert_frame_equal(df, df2)
         os.remove(path)