Implementation of Stata 13 (format 117) support

Author: PKEuS

pandas/io/stata.py

@@ -42,6 +42,7 @@ def read_stata(filepath_or_buffer, convert_dates=True, convert_categoricals=True
 
 _date_formats = ["%tc", "%tC", "%td", "%tw", "%tm", "%tq", "%th", "%ty"]
 
+
 def _stata_elapsed_date_to_datetime(date, fmt):
     """
     Convert from SIF to datetime. http://www.stata.com/help.cgi?datetime
@@ -234,7 +235,28 @@ def __init__(self, encoding=None):
                     (255, np.float64)
                 ]
             )
+        self.DTYPE_MAP_XML = \
+            dict(
+                [
+                    (32768, np.string_),
+                    (65526, np.float64),
+                    (65527, np.float32),
+                    (65528, np.int64),
+                    (65529, np.int32),
+                    (65530, np.int16)
+                ]
+            )
         self.TYPE_MAP = lrange(251) + list('bhlfd')
+        self.TYPE_MAP_XML = \
+            dict(
+                [
+                    (65526, 'd'),
+                    (65527, 'f'),
+                    (65528, 'l'),
+                    (65529, 'h'),
+                    (65530, 'b')
+                ]
+            )
         #NOTE: technically, some of these are wrong. there are more numbers
         # that can be represented. it's the 27 ABOVE and BELOW the max listed
         # numeric data type in [U] 12.2.2 of the 11.2 manual
@@ -304,86 +326,159 @@ def __init__(self, path_or_buf, encoding='cp1252'):
         self._read_header()
 
     def _read_header(self):
-        # header
-        self.format_version = struct.unpack('b', self.path_or_buf.read(1))[0]
-        if self.format_version not in [104, 105, 108, 113, 114, 115]:
-            raise ValueError("Version of given Stata file is not 104, 105, 108, 113 (Stata 8/9), 114 (Stata 10/11) or 115 (Stata 12)")
-        self.byteorder = self.path_or_buf.read(1) == 0x1 and '>' or '<'
-        self.filetype = struct.unpack('b', self.path_or_buf.read(1))[0]
-        self.path_or_buf.read(1)  # unused
-
-        self.nvar = struct.unpack(self.byteorder + 'H', self.path_or_buf.read(2))[0]
-        self.nobs = struct.unpack(self.byteorder + 'I', self.path_or_buf.read(4))[0]
-        if self.format_version > 105:
-            self.data_label = self.path_or_buf.read(81)
-        else:
-            self.data_label = self.path_or_buf.read(32)
-        if self.format_version > 104:
-            self.time_stamp = self.path_or_buf.read(18)
+        first_char = self.path_or_buf.read(1)
+        if struct.unpack('c', first_char)[0] is b'<':  # format 117 or higher (XML like)
+            self.path_or_buf.read(27)  # stata_dta><header><release>
+            self.format_version = int(self.path_or_buf.read(3))
+            if self.format_version not in [117]:
+                raise ValueError("Version of given Stata file is not 104, 105, 108, 113 (Stata 8/9), 114 (Stata 10/11), 115 (Stata 12) or 117 (Stata 13)")
+            self.path_or_buf.read(21)  # </release><byteorder>
+            self.byteorder = self.path_or_buf.read(3) == "LSF" and '>' or '<'
+            self.path_or_buf.read(15)  # </byteorder><K>
+            self.nvar = struct.unpack(self.byteorder + 'H', self.path_or_buf.read(2))[0]
+            self.path_or_buf.read(7)  # </K><N>
+            self.nobs = struct.unpack(self.byteorder + 'I', self.path_or_buf.read(4))[0]
+            self.path_or_buf.read(11)  # </N><label>
+            strlen = struct.unpack('b', self.path_or_buf.read(1))[0]
+            self.data_label = self.path_or_buf.read(strlen)
+            self.path_or_buf.read(19)  # </label><timestamp>
+            strlen = struct.unpack('b', self.path_or_buf.read(1))[0]
+            self.time_stamp = self.path_or_buf.read(strlen)
+            self.path_or_buf.read(26)  # </timestamp></header><map>
+            self.path_or_buf.read(8)  # 0x0000000000000000
+            self.path_or_buf.read(8)  # position of <map>
+            seek_vartypes = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 16
+            seek_varnames = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 10
+            seek_sortlist = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 10
+            seek_formats = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 9
+            seek_value_label_names = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 19
+            seek_variable_labels = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 17
+            self.path_or_buf.read(8)  # <characteristics>
+            self.data_location = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 6
+            self.seek_strls = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 7
+            self.seek_value_labels = struct.unpack(self.byteorder + 'q', self.path_or_buf.read(8))[0] + 14
+            #self.path_or_buf.read(8)  # </stata_dta>
+            #self.path_or_buf.read(8)  # EOF
+            self.path_or_buf.seek(seek_vartypes)
+            typlist = [struct.unpack(self.byteorder + 'H', self.path_or_buf.read(2))[0] for i in range(self.nvar)]
+            self.typlist = [None]*self.nvar
+            try:
+                i = 0
+                for typ in typlist:
+                    if typ <= 2045 or typ == 32768:
+                        self.typlist[i] = None
+                    else:
+                        self.typlist[i] = self.TYPE_MAP_XML[typ]
+                    i += 1
+            except:
+                raise ValueError("cannot convert stata types [{0}]".format(','.join(typlist)))
+            self.dtyplist = [None]*self.nvar
+            try:
+                i = 0
+                for typ in typlist:
+                    if typ <= 2045:
+                        self.dtyplist[i] = str(typ)
+                    else:
+                        self.dtyplist[i] = self.DTYPE_MAP_XML[typ]
+                    i += 1
+            except:
+                raise ValueError("cannot convert stata dtypes [{0}]".format(','.join(typlist)))
 
-        # descriptors
-        if self.format_version > 108:
-            typlist = [ord(self.path_or_buf.read(1)) for i in range(self.nvar)]
-        else:
-            typlist = [self.OLD_TYPE_MAPPING[self._decode_bytes(self.path_or_buf.read(1))] for i in range(self.nvar)]
-
-        try:
-            self.typlist = [self.TYPE_MAP[typ] for typ in typlist]
-        except:
-            raise ValueError("cannot convert stata types [{0}]".format(','.join(typlist)))
-        try:
-            self.dtyplist = [self.DTYPE_MAP[typ] for typ in typlist]
-        except:
-            raise ValueError("cannot convert stata dtypes [{0}]".format(','.join(typlist)))
-
-        if self.format_version > 108:
+            self.path_or_buf.seek(seek_varnames)
             self.varlist = [self._null_terminate(self.path_or_buf.read(33)) for i in range(self.nvar)]
-        else:
-            self.varlist = [self._null_terminate(self.path_or_buf.read(9)) for i in range(self.nvar)]
-        self.srtlist = struct.unpack(self.byteorder + ('h' * (self.nvar + 1)), self.path_or_buf.read(2 * (self.nvar + 1)))[:-1]
-        if self.format_version > 113:
+
+            self.path_or_buf.seek(seek_sortlist)
+            self.srtlist = struct.unpack(self.byteorder + ('h' * (self.nvar + 1)), self.path_or_buf.read(2 * (self.nvar + 1)))[:-1]
+
+            self.path_or_buf.seek(seek_formats)
             self.fmtlist = [self._null_terminate(self.path_or_buf.read(49)) for i in range(self.nvar)]
-        elif self.format_version > 104:
-            self.fmtlist = [self._null_terminate(self.path_or_buf.read(12)) for i in range(self.nvar)]
-        else:
-            self.fmtlist = [self._null_terminate(self.path_or_buf.read(7)) for i in range(self.nvar)]
-        if self.format_version > 108:
+
+            self.path_or_buf.seek(seek_value_label_names)
             self.lbllist = [self._null_terminate(self.path_or_buf.read(33)) for i in range(self.nvar)]
-        else:
-            self.lbllist = [self._null_terminate(self.path_or_buf.read(9)) for i in range(self.nvar)]
-        if self.format_version > 105:
+
+            self.path_or_buf.seek(seek_variable_labels)
             self.vlblist = [self._null_terminate(self.path_or_buf.read(81)) for i in range(self.nvar)]
         else:
-            self.vlblist = [self._null_terminate(self.path_or_buf.read(32)) for i in range(self.nvar)]
-
-        # ignore expansion fields (Format 105 and later)
-        # When reading, read five bytes; the last four bytes now tell you the
-        # size of the next read, which you discard.  You then continue like
-        # this until you read 5 bytes of zeros.
-
-        if self.format_version > 104:
-            while True:
-                data_type = struct.unpack(self.byteorder + 'b', self.path_or_buf.read(1))[0]
-                if self.format_version > 108:
-                    data_len = struct.unpack(self.byteorder + 'i', self.path_or_buf.read(4))[0]
-                else:
-                    data_len = struct.unpack(self.byteorder + 'h', self.path_or_buf.read(2))[0]
-                if data_type == 0:
-                    break
-                self.path_or_buf.read(data_len)
+            # header
+            self.format_version = struct.unpack('b', first_char)[0]
+            if self.format_version not in [104, 105, 108, 113, 114, 115]:
+                raise ValueError("Version of given Stata file is not 104, 105, 108, 113 (Stata 8/9), 114 (Stata 10/11), 115 (Stata 12) or 117 (Stata 13)")
+            self.byteorder = self.path_or_buf.read(1) == 0x1 and '>' or '<'
+            self.filetype = struct.unpack('b', self.path_or_buf.read(1))[0]
+            self.path_or_buf.read(1)  # unused
+
+            self.nvar = struct.unpack(self.byteorder + 'H', self.path_or_buf.read(2))[0]
+            self.nobs = struct.unpack(self.byteorder + 'I', self.path_or_buf.read(4))[0]
+            if self.format_version > 105:
+                self.data_label = self.path_or_buf.read(81)
+            else:
+                self.data_label = self.path_or_buf.read(32)
+            if self.format_version > 104:
+                self.time_stamp = self.path_or_buf.read(18)
+
+            # descriptors
+            if self.format_version > 108:
+                typlist = [ord(self.path_or_buf.read(1)) for i in range(self.nvar)]
+            else:
+                typlist = [self.OLD_TYPE_MAPPING[self._decode_bytes(self.path_or_buf.read(1))] for i in range(self.nvar)]
+
+            try:
+                self.typlist = [self.TYPE_MAP[typ] for typ in typlist]
+            except:
+                raise ValueError("cannot convert stata types [{0}]".format(','.join(typlist)))
+            try:
+                self.dtyplist = [self.DTYPE_MAP[typ] for typ in typlist]
+            except:
+                raise ValueError("cannot convert stata dtypes [{0}]".format(','.join(typlist)))
+
+            if self.format_version > 108:
+                self.varlist = [self._null_terminate(self.path_or_buf.read(33)) for i in range(self.nvar)]
+            else:
+                self.varlist = [self._null_terminate(self.path_or_buf.read(9)) for i in range(self.nvar)]
+            self.srtlist = struct.unpack(self.byteorder + ('h' * (self.nvar + 1)), self.path_or_buf.read(2 * (self.nvar + 1)))[:-1]
+            if self.format_version > 113:
+                self.fmtlist = [self._null_terminate(self.path_or_buf.read(49)) for i in range(self.nvar)]
+            elif self.format_version > 104:
+                self.fmtlist = [self._null_terminate(self.path_or_buf.read(12)) for i in range(self.nvar)]
+            else:
+                self.fmtlist = [self._null_terminate(self.path_or_buf.read(7)) for i in range(self.nvar)]
+            if self.format_version > 108:
+                self.lbllist = [self._null_terminate(self.path_or_buf.read(33)) for i in range(self.nvar)]
+            else:
+                self.lbllist = [self._null_terminate(self.path_or_buf.read(9)) for i in range(self.nvar)]
+            if self.format_version > 105:
+                self.vlblist = [self._null_terminate(self.path_or_buf.read(81)) for i in range(self.nvar)]
+            else:
+                self.vlblist = [self._null_terminate(self.path_or_buf.read(32)) for i in range(self.nvar)]
+
+            # ignore expansion fields (Format 105 and later)
+            # When reading, read five bytes; the last four bytes now tell you the
+            # size of the next read, which you discard.  You then continue like
+            # this until you read 5 bytes of zeros.
+
+            if self.format_version > 104:
+                while True:
+                    data_type = struct.unpack(self.byteorder + 'b', self.path_or_buf.read(1))[0]
+                    if self.format_version > 108:
+                        data_len = struct.unpack(self.byteorder + 'i', self.path_or_buf.read(4))[0]
+                    else:
+                        data_len = struct.unpack(self.byteorder + 'h', self.path_or_buf.read(2))[0]
+                    if data_type == 0:
+                        break
+                    self.path_or_buf.read(data_len)
+
+            # necessary data to continue parsing
+            self.data_location = self.path_or_buf.tell()
 
-        # necessary data to continue parsing
-        self.data_location = self.path_or_buf.tell()
         self.has_string_data = len([x for x in self.typlist if type(x) is int]) > 0
-        self._col_size()
+
+        """Calculate size of a data record."""
+        self.col_sizes = lmap(lambda x: self._calcsize(x), self.typlist)
 
     def _calcsize(self, fmt):
         return type(fmt) is int and fmt or struct.calcsize(self.byteorder + fmt)
 
     def _col_size(self, k=None):
-        """Calculate size of a data record."""
-        if len(self.col_sizes) == 0:
-            self.col_sizes = lmap(lambda x: self._calcsize(x), self.typlist)
         if k is None:
             return self.col_sizes
         else:
@@ -427,8 +522,8 @@ def _next(self):
             return data
         else:
             return list(map(lambda i: self._unpack(typlist[i],
-                                              self.path_or_buf.read(self._col_size(i))),
-                       range(self.nvar)))
+                                                   self.path_or_buf.read(self._col_size(i))),
+                            range(self.nvar)))
 
     def _dataset(self):
         """
@@ -450,29 +545,33 @@ def _dataset(self):
         be handled by your applcation.
         """
 
-        try:
-            self._file.seek(self._data_location)
-        except Exception:
-            pass
+        self.path_or_buf.seek(self.data_location)
 
         for i in range(self.nobs):
             yield self._next()
 
     def _read_value_labels(self):
-        if not self._data_read:
-            raise Exception("Data has not been read. Because of the layout of Stata files, this is necessary before reading value labels.")
-        if self._value_labels_read:
-            raise Exception("Value labels have already been read.")
+        if self.format_version >= 117:
+            self.path_or_buf.seek(self.seek_value_labels)
+        else:
+            if not self._data_read:
+                raise Exception("Data has not been read. Because of the layout of Stata files, this is necessary before reading value labels.")
+            if self._value_labels_read:
+                raise Exception("Value labels have already been read.")
 
         self.value_label_dict = dict()
 
         if self.format_version <= 108:
             return  # Value labels are not supported in version 108 and earlier.
 
         while True:
+            if self.format_version >= 117:
+                if self._decode_bytes(self.path_or_buf.read(5), self._encoding) == '</val':  # <lbl>
+                    break  # end o f variable lable table
+
             slength = self.path_or_buf.read(4)
             if not slength:
-                break  # end of variable lable table
+                break  # end of variable lable table (format < 117)
             labname = self._null_terminate(self.path_or_buf.read(33))
             self.path_or_buf.read(3)  # padding
 
@@ -488,8 +587,24 @@ def _read_value_labels(self):
             self.value_label_dict[labname] = dict()
             for i in range(n):
                 self.value_label_dict[labname][val[i]] = self._null_terminate(txt[off[i]:])
+
+            if self.format_version >= 117:
+                self.path_or_buf.read(6)  # </lbl>
         self._value_labels_read = True
 
+    def _read_strls(self):
+        self.path_or_buf.seek(self.seek_strls)
+        self.GSO = dict()
+        while True:
+            if self.path_or_buf.read(3) is not 'GSO':
+                break
+
+            v_o = struct.unpack(self.byteorder + 'L', self.path_or_buf.read(8))[0]
+            typ = self.path_or_buf.read(1)
+            length = struct.unpack(self.byteorder + 'I', self.path_or_buf.read(4))[0]
+            self.GSO[v_o] = self.path_or_buf.read(length-1)
+            self.path_or_buf.read(1)  # zero-termination
+
     def data(self, convert_dates=True, convert_categoricals=True, index=None):
         """
         Reads observations from Stata file, converting them into a dataframe
@@ -511,6 +626,9 @@ def data(self, convert_dates=True, convert_categoricals=True, index=None):
             raise Exception("Data has already been read.")
         self._data_read = True
 
+        if self.format_version >= 117:
+            self._read_strls()
+
         stata_dta = self._dataset()
 
         data = []

pandas/io/tests/test_stata.py

@@ -33,7 +33,6 @@ def setUp(self):
         self.dta9 = os.path.join(self.dirpath, 'lbw.dta')
         self.csv9 = os.path.join(self.dirpath, 'lbw.csv')
         self.dta_encoding = os.path.join(self.dirpath, 'stata1_encoding.dta')
-
     def read_dta(self, file):
         return read_stata(file, convert_dates=True)
 
@@ -198,6 +197,17 @@ def test_read_write_dta10(self):
             tm.assert_frame_equal(written_and_read_again.set_index('index'),
                                   original)
 
+    def test_read_dta11(self):
+        reader = StataReader(self.dta11)
+        parsed = reader.data()
+        # Pandas uses np.nan as missing value.
+        # Thus, all columns will be of type float, regardless of their name.
+        expected = DataFrame([(np.nan, np.nan, np.nan, np.nan, np.nan)],
+                             columns=['float_miss', 'double_miss', 'byte_miss',
+                                      'int_miss', 'long_miss'])
+
+        tm.assert_frame_equal(parsed, expected)
+
     def test_stata_doc_examples(self):
         with tm.ensure_clean() as path:
             df = DataFrame(np.random.randn(10, 2), columns=list('AB'))