BUG: Enforce correct encoding in stata

Ensure StataReader and StataWriter have the correct encoding.
Standardized default encoding to 'latin-1'

closes pandas-dev#15723

Author: Kevin Sheppard <[email protected]>

Closes pandas-dev#15768 from bashtage/limit-stata-encoding and squashes the following commits:

8278be7 [Kevin Sheppard] BUG: Fix limited key range on 32-bit platofrms
2f02697 [Kevin Sheppard] BUG: Enforce correct encoding in stata

Author: bashtage

doc/source/whatsnew/v0.20.0.txt

@@ -919,6 +919,8 @@ Bug Fixes
 - Avoid use of ``np.finfo()`` during ``import pandas`` removed to mitigate deadlock on Python GIL misuse (:issue:`14641`)
 
 - Bug in ``DataFrame.to_stata()`` and ``StataWriter`` which produces incorrectly formatted files to be produced for some locales (:issue:`13856`)
+- Bug in ``StataReader`` and ``StataWriter`` which allows invalid encodings (:issue:`15723`)
+
 - Bug in ``pd.concat()`` in which concatting with an empty dataframe with ``join='inner'`` was being improperly handled (:issue:`15328`)
 - Bug in ``groupby.agg()`` incorrectly localizing timezone on ``datetime`` (:issue:`15426`, :issue:`10668`, :issue:`13046`)
 
@@ -933,3 +935,4 @@ Bug Fixes
 - Bug in ``pd.melt()`` where passing a tuple value for ``value_vars`` caused a ``TypeError`` (:issue:`15348`)
 - Bug in ``.eval()`` which caused multiline evals to fail with local variables not on the first line (:issue:`15342`)
 - Bug in ``pd.read_msgpack`` which did not allow to load dataframe with an index of type ``CategoricalIndex`` (:issue:`15487`)
+

pandas/io/stata.py

@@ -33,6 +33,9 @@
 from pandas._libs.lib import max_len_string_array, infer_dtype
 from pandas._libs.tslib import NaT, Timestamp
 
+VALID_ENCODINGS = ('ascii', 'us-ascii', 'latin-1', 'latin_1', 'iso-8859-1',
+                   'iso8859-1', '8859', 'cp819', 'latin', 'latin1', 'L1')
+
 _version_error = ("Version of given Stata file is not 104, 105, 108, "
                   "111 (Stata 7SE), 113 (Stata 8/9), 114 (Stata 10/11), "
                   "115 (Stata 12), 117 (Stata 13), or 118 (Stata 14)")
@@ -45,7 +48,7 @@
 
 _encoding_params = """\
 encoding : string, None or encoding
-    Encoding used to parse the files. None defaults to iso-8859-1."""
+    Encoding used to parse the files. None defaults to latin-1."""
 
 _statafile_processing_params2 = """\
 index : identifier of index column
@@ -816,9 +819,14 @@ def get_base_missing_value(cls, dtype):
 
 
 class StataParser(object):
-    _default_encoding = 'iso-8859-1'
+    _default_encoding = 'latin-1'
 
     def __init__(self, encoding):
+        if encoding is not None:
+            if encoding not in VALID_ENCODINGS:
+                raise ValueError('Unknown encoding. Only latin-1 and ascii '
+                                 'supported.')
+
         self._encoding = encoding
 
         # type          code.
@@ -936,7 +944,7 @@ def __init__(self, path_or_buf, convert_dates=True,
                  convert_categoricals=True, index=None,
                  convert_missing=False, preserve_dtypes=True,
                  columns=None, order_categoricals=True,
-                 encoding='iso-8859-1', chunksize=None):
+                 encoding='latin-1', chunksize=None):
         super(StataReader, self).__init__(encoding)
         self.col_sizes = ()
 
@@ -949,6 +957,10 @@ def __init__(self, path_or_buf, convert_dates=True,
         self._preserve_dtypes = preserve_dtypes
         self._columns = columns
         self._order_categoricals = order_categoricals
+        if encoding is not None:
+            if encoding not in VALID_ENCODINGS:
+                raise ValueError('Unknown encoding. Only latin-1 and  ascii '
+                                 'supported.')
         self._encoding = encoding
         self._chunksize = chunksize
 
@@ -1362,7 +1374,8 @@ def _read_value_labels(self):
 
     def _read_strls(self):
         self.path_or_buf.seek(self.seek_strls)
-        self.GSO = {0: ''}
+        # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
+        self.GSO = {'0': ''}
         while True:
             if self.path_or_buf.read(3) != b'GSO':
                 break
@@ -1387,7 +1400,8 @@ def _read_strls(self):
                 if self.format_version == 117:
                     encoding = self._encoding or self._default_encoding
                 va = va[0:-1].decode(encoding)
-            self.GSO[v_o] = va
+            # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
+            self.GSO[str(v_o)] = va
 
     # legacy
     @Appender('DEPRECATED: ' + _data_method_doc)
@@ -1623,7 +1637,8 @@ def _insert_strls(self, data):
         for i, typ in enumerate(self.typlist):
             if typ != 'Q':
                 continue
-            data.iloc[:, i] = [self.GSO[k] for k in data.iloc[:, i]]
+            # Wrap v_o in a string to allow uint64 values as keys on 32bit OS
+            data.iloc[:, i] = [self.GSO[str(k)] for k in data.iloc[:, i]]
         return data
 
     def _do_select_columns(self, data, columns):
@@ -1855,7 +1870,7 @@ class StataWriter(StataParser):
     write_index : bool
         Write the index to Stata dataset.
     encoding : str
-        Default is latin-1. Unicode is not supported
+        Default is latin-1. Only latin-1 and ascii are supported.
     byteorder : str
         Can be ">", "<", "little", or "big". default is `sys.byteorder`
     time_stamp : datetime

pandas/tests/io/test_stata.py

@@ -1276,3 +1276,10 @@ def test_out_of_range_float(self):
                 original.to_stata(path)
             tm.assertTrue('ColumnTooBig' in cm.exception)
             tm.assertTrue('infinity' in cm.exception)
+
+    def test_invalid_encoding(self):
+        # GH15723, validate encoding
+        original = self.read_csv(self.csv3)
+        with tm.assertRaises(ValueError):
+            with tm.ensure_clean() as path:
+                original.to_stata(path, encoding='utf-8')