<s> """<STR_LIT>""" <EOL> from statsmodels . compat . python import ( zip , lzip , lmap , lrange , string_types , long , lfilter , <EOL> asbytes , asstr , range , PY3 ) <EOL> from struct import unpack , calcsize , pack <EOL> from struct import error as struct_error <EOL> import datetime <EOL> import sys <EOL> import numpy as np <EOL> from numpy . lib . _iotools import _is_string_like , easy_dtype <EOL> import statsmodels . tools . data as data_util <EOL> from pandas import isnull <EOL> from statsmodels . iolib . openfile import get_file_obj <EOL> _date_formats = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> def _datetime_to_stata_elapsed ( date , fmt ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( date , datetime . datetime ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> stata_epoch = datetime . datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> if fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> delta = date - stata_epoch <EOL> return ( delta . days * <NUM_LIT> + delta . seconds * <NUM_LIT:1000> + <EOL> delta . microseconds / <NUM_LIT:1000> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> from warnings import warn <EOL> warn ( "<STR_LIT>" , UserWarning ) <EOL> return date <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return ( date - stata_epoch ) . days <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return ( <NUM_LIT> * ( date . year - stata_epoch . year ) + <EOL> ( date - datetime . datetime ( date . year , <NUM_LIT:1> , <NUM_LIT:1> ) ) . days / <NUM_LIT:7> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return ( <NUM_LIT:12> * ( date . year - stata_epoch . year ) + date . month - <NUM_LIT:1> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return <NUM_LIT:4> * ( date . year - stata_epoch . year ) + int ( ( date . month - <NUM_LIT:1> ) / <NUM_LIT:3> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return <NUM_LIT:2> * ( date . year - stata_epoch . year ) + int ( date . month > <NUM_LIT:6> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return date . year <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % fmt ) <EOL> def _stata_elapsed_date_to_datetime ( date , fmt ) : <EOL> """<STR_LIT>""" <EOL> date = int ( date ) <EOL> stata_epoch = datetime . datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> if fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> from dateutil . relativedelta import relativedelta <EOL> return stata_epoch + relativedelta ( microseconds = date * <NUM_LIT:1000> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> from warnings import warn <EOL> warn ( "<STR_LIT>" , <EOL> UserWarning ) <EOL> return date <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return stata_epoch + datetime . timedelta ( int ( date ) ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> year = datetime . datetime ( stata_epoch . year + date // <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> day_delta = ( date % <NUM_LIT> ) * <NUM_LIT:7> <EOL> return year + datetime . timedelta ( int ( day_delta ) ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> year = stata_epoch . year + date // <NUM_LIT:12> <EOL> month_delta = ( date % <NUM_LIT:12> ) + <NUM_LIT:1> <EOL> return datetime . datetime ( year , month_delta , <NUM_LIT:1> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> year = stata_epoch . year + date // <NUM_LIT:4> <EOL> month_delta = ( date % <NUM_LIT:4> ) * <NUM_LIT:3> + <NUM_LIT:1> <EOL> return datetime . datetime ( year , month_delta , <NUM_LIT:1> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> year = stata_epoch . year + date // <NUM_LIT:2> <EOL> month_delta = ( date % <NUM_LIT:2> ) * <NUM_LIT:6> + <NUM_LIT:1> <EOL> return datetime . datetime ( year , month_delta , <NUM_LIT:1> ) <EOL> elif fmt in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> if date > <NUM_LIT:0> : <EOL> return datetime . datetime ( date , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % fmt ) <EOL> class _StataMissingValue ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , offset , value ) : <EOL> self . _value = value <EOL> if isinstance ( value , ( int , long ) ) : <EOL> self . _str = value - offset is <NUM_LIT:1> and '<STR_LIT:.>' or ( '<STR_LIT:.>' + chr ( value - offset + <NUM_LIT> ) ) <EOL> else : <EOL> self . _str = '<STR_LIT:.>' <EOL> string = property ( lambda self : self . _str , doc = "<STR_LIT>" ) <EOL> value = property ( lambda self : self . _value , doc = '<STR_LIT>' ) <EOL> def __str__ ( self ) : return self . _str <EOL> __str__ . __doc__ = string . __doc__ <EOL> class _StataVariable ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , variable_data ) : <EOL> self . _data = variable_data <EOL> def __int__ ( self ) : <EOL> return self . index <EOL> def __str__ ( self ) : <EOL> return self . name <EOL> index = property ( lambda self : self . _data [ <NUM_LIT:0> ] , doc = '<STR_LIT>' ) <EOL> type = property ( lambda self : self . _data [ <NUM_LIT:1> ] , doc = '<STR_LIT>' ) <EOL> name = property ( lambda self : self . _data [ <NUM_LIT:2> ] , doc = '<STR_LIT>' ) <EOL> format = property ( lambda self : self . _data [ <NUM_LIT:4> ] , doc = '<STR_LIT>' ) <EOL> value_format = property ( lambda self : self . _data [ <NUM_LIT:5> ] , doc = '<STR_LIT>' ) <EOL> label = property ( lambda self : self . _data [ <NUM_LIT:6> ] , doc = '<STR_LIT>' ) <EOL> __int__ . __doc__ = index . __doc__ <EOL> __str__ . __doc__ = name . __doc__ <EOL> class StataReader ( object ) : <EOL> """<STR_LIT>""" <EOL> _header = { } <EOL> _data_location = <NUM_LIT:0> <EOL> _col_sizes = ( ) <EOL> _has_string_data = False <EOL> _missing_values = False <EOL> DTYPE_MAP = dict ( lzip ( lrange ( <NUM_LIT:1> , <NUM_LIT> ) , [ '<STR_LIT:a>' + str ( i ) for i in range ( <NUM_LIT:1> , <NUM_LIT> ) ] ) + [ ( <NUM_LIT> , np . int16 ) , ( <NUM_LIT> , np . int32 ) , ( <NUM_LIT> , int ) , <EOL> ( <NUM_LIT> , np . float32 ) , ( <NUM_LIT:255> , np . float64 ) ] ) <EOL> TYPE_MAP = lrange ( <NUM_LIT> ) + list ( '<STR_LIT>' ) <EOL> MISSING_VALUES = { '<STR_LIT:b>' : ( - <NUM_LIT> , <NUM_LIT:100> ) , '<STR_LIT:h>' : ( - <NUM_LIT> , <NUM_LIT> ) , '<STR_LIT:l>' : <EOL> ( - <NUM_LIT> , <NUM_LIT> ) , '<STR_LIT:f>' : ( - <NUM_LIT> , + <NUM_LIT> ) , '<STR_LIT:d>' : <EOL> ( - <NUM_LIT> , + <NUM_LIT> ) } <EOL> def __init__ ( self , fname , missing_values = False , encoding = None ) : <EOL> if encoding == None : <EOL> import locale <EOL> self . _encoding = locale . getpreferredencoding ( ) <EOL> else : <EOL> self . _encoding = encoding <EOL> self . _missing_values = missing_values <EOL> self . _parse_header ( fname ) <EOL> def file_headers ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _header <EOL> def file_format ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _header [ '<STR_LIT>' ] <EOL> def file_label ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _header [ '<STR_LIT>' ] <EOL> def file_timestamp ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _header [ '<STR_LIT>' ] <EOL> def variables ( self ) : <EOL> """<STR_LIT>""" <EOL> return lmap ( _StataVariable , zip ( lrange ( self . _header [ '<STR_LIT>' ] ) , <EOL> self . _header [ '<STR_LIT>' ] , self . _header [ '<STR_LIT>' ] , <EOL> self . _header [ '<STR_LIT>' ] , <EOL> self . _header [ '<STR_LIT>' ] , self . _header [ '<STR_LIT>' ] , <EOL> self . _header [ '<STR_LIT>' ] ) ) <EOL> def dataset ( self , as_dict = False ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> self . _file . seek ( self . _data_location ) <EOL> except Exception : <EOL> pass <EOL> if as_dict : <EOL> vars = lmap ( str , self . variables ( ) ) <EOL> for i in range ( len ( self ) ) : <EOL> yield dict ( zip ( vars , self . _next ( ) ) ) <EOL> else : <EOL> for i in range ( self . _header [ '<STR_LIT>' ] ) : <EOL> yield self . _next ( ) <EOL> def __len__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _header [ '<STR_LIT>' ] <EOL> def __getitem__ ( self , k ) : <EOL> """<STR_LIT>""" <EOL> if not ( isinstance ( k , ( int , long ) ) ) or k < <NUM_LIT:0> or k > len ( self ) - <NUM_LIT:1> : <EOL> raise IndexError ( k ) <EOL> loc = self . _data_location + sum ( self . _col_size ( ) ) * k <EOL> if self . _file . tell ( ) != loc : <EOL> self . _file . seek ( loc ) <EOL> return self . _next ( ) <EOL> def _null_terminate ( self , s , encoding ) : <EOL> if PY3 : <EOL> null_byte = asbytes ( '<STR_LIT:\x00>' ) <EOL> try : <EOL> s = s . lstrip ( null_byte ) [ : s . index ( null_byte ) ] <EOL> except : <EOL> pass <EOL> return s . decode ( encoding ) <EOL> else : <EOL> null_byte = asbytes ( '<STR_LIT:\x00>' ) <EOL> try : <EOL> return s . lstrip ( null_byte ) [ : s . index ( null_byte ) ] <EOL> except : <EOL> return s <EOL> def _parse_header ( self , file_object ) : <EOL> self . _file = file_object <EOL> encoding = self . _encoding <EOL> self . _header [ '<STR_LIT>' ] = unpack ( '<STR_LIT:b>' , self . _file . read ( <NUM_LIT:1> ) ) [ <NUM_LIT:0> ] <EOL> if self . _header [ '<STR_LIT>' ] not in [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> self . _header [ '<STR_LIT>' ] ) <EOL> byteorder = self . _header [ '<STR_LIT>' ] = unpack ( '<STR_LIT:b>' , <EOL> self . _file . read ( <NUM_LIT:1> ) ) [ <NUM_LIT:0> ] == <NUM_LIT> and '<STR_LIT:>>' or '<STR_LIT:<>' <EOL> self . _header [ '<STR_LIT>' ] = unpack ( '<STR_LIT:b>' , self . _file . read ( <NUM_LIT:1> ) ) [ <NUM_LIT:0> ] <EOL> self . _file . read ( <NUM_LIT:1> ) <EOL> nvar = self . _header [ '<STR_LIT>' ] = unpack ( byteorder + '<STR_LIT:h>' , <EOL> self . _file . read ( <NUM_LIT:2> ) ) [ <NUM_LIT:0> ] <EOL> self . _header [ '<STR_LIT>' ] = unpack ( byteorder + '<STR_LIT:i>' , self . _file . read ( <NUM_LIT:4> ) ) [ <NUM_LIT:0> ] <EOL> self . _header [ '<STR_LIT>' ] = self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , <EOL> encoding ) <EOL> self . _header [ '<STR_LIT>' ] = self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , <EOL> encoding ) <EOL> typlist = [ ord ( self . _file . read ( <NUM_LIT:1> ) ) for i in range ( nvar ) ] <EOL> self . _header [ '<STR_LIT>' ] = [ self . TYPE_MAP [ typ ] for typ in typlist ] <EOL> self . _header [ '<STR_LIT>' ] = [ self . DTYPE_MAP [ typ ] for typ in typlist ] <EOL> self . _header [ '<STR_LIT>' ] = [ self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , <EOL> encoding ) for i in range ( nvar ) ] <EOL> self . _header [ '<STR_LIT>' ] = unpack ( byteorder + ( '<STR_LIT:h>' * ( nvar + <NUM_LIT:1> ) ) , <EOL> self . _file . read ( <NUM_LIT:2> * ( nvar + <NUM_LIT:1> ) ) ) [ : - <NUM_LIT:1> ] <EOL> if self . _header [ '<STR_LIT>' ] <= <NUM_LIT> : <EOL> self . _header [ '<STR_LIT>' ] = [ self . _null_terminate ( self . _file . read ( <NUM_LIT:12> ) , encoding ) for i in range ( nvar ) ] <EOL> else : <EOL> self . _header [ '<STR_LIT>' ] = [ self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , encoding ) for i in range ( nvar ) ] <EOL> self . _header [ '<STR_LIT>' ] = [ self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , <EOL> encoding ) for i in range ( nvar ) ] <EOL> self . _header [ '<STR_LIT>' ] = [ self . _null_terminate ( self . _file . read ( <NUM_LIT> ) , <EOL> encoding ) for i in range ( nvar ) ] <EOL> while True : <EOL> data_type = unpack ( byteorder + '<STR_LIT:b>' , self . _file . read ( <NUM_LIT:1> ) ) [ <NUM_LIT:0> ] <EOL> data_len = unpack ( byteorder + '<STR_LIT:i>' , self . _file . read ( <NUM_LIT:4> ) ) [ <NUM_LIT:0> ] <EOL> if data_type == <NUM_LIT:0> : <EOL> break <EOL> self . _file . read ( data_len ) <EOL> self . _data_location = self . _file . tell ( ) <EOL> self . _has_string_data = len ( lfilter ( lambda x : isinstance ( x , int ) , <EOL> self . _header [ '<STR_LIT>' ] ) ) > <NUM_LIT:0> <EOL> self . _col_size ( ) <EOL> def _calcsize ( self , fmt ) : <EOL> return isinstance ( fmt , int ) and fmt or calcsize ( self . _header [ '<STR_LIT>' ] + fmt ) <EOL> def _col_size ( self , k = None ) : <EOL> """<STR_LIT>""" <EOL> if len ( self . _col_sizes ) == <NUM_LIT:0> : <EOL> self . _col_sizes = lmap ( lambda x : self . _calcsize ( x ) , <EOL> self . _header [ '<STR_LIT>' ] ) <EOL> if k == None : <EOL> return self . _col_sizes <EOL> else : <EOL> return self . _col_sizes [ k ] <EOL> def _unpack ( self , fmt , byt ) : <EOL> d = unpack ( self . _header [ '<STR_LIT>' ] + fmt , byt ) [ <NUM_LIT:0> ] <EOL> if fmt [ - <NUM_LIT:1> ] in self . MISSING_VALUES : <EOL> nmin , nmax = self . MISSING_VALUES [ fmt [ - <NUM_LIT:1> ] ] <EOL> if d < nmin or d > nmax : <EOL> if self . _missing_values : <EOL> return _StataMissingValue ( nmax , d ) <EOL> else : <EOL> return None <EOL> return d <EOL> def _next ( self ) : <EOL> typlist = self . _header [ '<STR_LIT>' ] <EOL> if self . _has_string_data : <EOL> data = [ None ] * self . _header [ '<STR_LIT>' ] <EOL> for i in range ( len ( data ) ) : <EOL> if isinstance ( typlist [ i ] , int ) : <EOL> data [ i ] = self . _null_terminate ( self . _file . read ( typlist [ i ] ) , <EOL> self . _encoding ) <EOL> else : <EOL> data [ i ] = self . _unpack ( typlist [ i ] , <EOL> self . _file . read ( self . _col_size ( i ) ) ) <EOL> return data <EOL> else : <EOL> return lmap ( lambda i : self . _unpack ( typlist [ i ] , <EOL> self . _file . read ( self . _col_size ( i ) ) ) , <EOL> lrange ( self . _header [ '<STR_LIT>' ] ) ) <EOL> def _set_endianness ( endianness ) : <EOL> if endianness . lower ( ) in [ "<STR_LIT:<>" , "<STR_LIT>" ] : <EOL> return "<STR_LIT:<>" <EOL> elif endianness . lower ( ) in [ "<STR_LIT:>>" , "<STR_LIT>" ] : <EOL> return "<STR_LIT:>>" <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % endianness ) <EOL> def _dtype_to_stata_type ( dtype ) : <EOL> """<STR_LIT>""" <EOL> if dtype . type == np . string_ : <EOL> return chr ( dtype . itemsize ) <EOL> elif dtype . type == np . object_ : <EOL> return chr ( <NUM_LIT> ) <EOL> elif dtype == np . float64 : <EOL> return chr ( <NUM_LIT:255> ) <EOL> elif dtype == np . float32 : <EOL> return chr ( <NUM_LIT> ) <EOL> elif dtype == np . int64 : <EOL> return chr ( <NUM_LIT> ) <EOL> elif dtype == np . int32 : <EOL> return chr ( <NUM_LIT> ) <EOL> elif dtype == np . int8 or dtype == np . int16 : <EOL> return chr ( <NUM_LIT> ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % dtype ) <EOL> def _dtype_to_default_stata_fmt ( dtype ) : <EOL> """<STR_LIT>""" <EOL> if dtype . type == np . string_ : <EOL> return "<STR_LIT:%>" + str ( dtype . itemsize ) + "<STR_LIT:s>" <EOL> elif dtype . type == np . object_ : <EOL> return "<STR_LIT>" <EOL> elif dtype == np . float64 : <EOL> return "<STR_LIT>" <EOL> elif dtype == np . float32 : <EOL> return "<STR_LIT>" <EOL> elif dtype == np . int64 : <EOL> return "<STR_LIT>" <EOL> elif dtype == np . int32 : <EOL> return "<STR_LIT>" <EOL> elif dtype == np . int8 or dtype == np . int16 : <EOL> return "<STR_LIT>" <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % dtype ) <EOL> def _pad_bytes ( name , length ) : <EOL> """<STR_LIT>""" <EOL> return name + "<STR_LIT:\x00>" * ( length - len ( name ) ) <EOL> def _default_names ( nvar ) : <EOL> """<STR_LIT>""" <EOL> return [ "<STR_LIT>" % i for i in range ( <NUM_LIT:1> , nvar + <NUM_LIT:1> ) ] <EOL> def _convert_datetime_to_stata_type ( fmt ) : <EOL> """<STR_LIT>""" <EOL> if fmt in [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> return np . float64 <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % fmt ) <EOL> def _maybe_convert_to_int_keys ( convert_dates , varlist ) : <EOL> new_dict = { } <EOL> for key in convert_dates : <EOL> if not convert_dates [ key ] . startswith ( "<STR_LIT:%>" ) : <EOL> convert_dates [ key ] = "<STR_LIT:%>" + convert_dates [ key ] <EOL> if key in varlist : <EOL> new_dict . update ( { varlist . index ( key ) : convert_dates [ key ] } ) <EOL> else : <EOL> if not isinstance ( key , int ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> new_dict . update ( { key : convert_dates [ key ] } ) <EOL> return new_dict <EOL> _type_converters = { <NUM_LIT> : np . long , <NUM_LIT> : int } <EOL> class StataWriter ( object ) : <EOL> """<STR_LIT>""" <EOL> DTYPE_MAP = dict ( lzip ( lrange ( <NUM_LIT:1> , <NUM_LIT> ) , [ '<STR_LIT:a>' + str ( i ) for i in range ( <NUM_LIT:1> , <NUM_LIT> ) ] ) + [ ( <NUM_LIT> , np . int16 ) , ( <NUM_LIT> , np . int32 ) , ( <NUM_LIT> , int ) , <EOL> ( <NUM_LIT> , np . float32 ) , ( <NUM_LIT:255> , np . float64 ) ] ) <EOL> TYPE_MAP = lrange ( <NUM_LIT> ) + list ( '<STR_LIT>' ) <EOL> MISSING_VALUES = { '<STR_LIT:b>' : <NUM_LIT> , <EOL> '<STR_LIT:h>' : <NUM_LIT> , <EOL> '<STR_LIT:l>' : <NUM_LIT> , <EOL> '<STR_LIT:f>' : <NUM_LIT> , <EOL> '<STR_LIT:d>' : <NUM_LIT> } <EOL> def __init__ ( self , fname , data , convert_dates = None , encoding = "<STR_LIT>" , <EOL> byteorder = None ) : <EOL> self . _convert_dates = convert_dates <EOL> if data_util . _is_using_pandas ( data , None ) : <EOL> self . _prepare_pandas ( data ) <EOL> elif data_util . _is_array_like ( data , None ) : <EOL> data = np . asarray ( data ) <EOL> if data_util . _is_structured_ndarray ( data ) : <EOL> self . _prepare_structured_array ( data ) <EOL> else : <EOL> if convert_dates is not None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> self . _prepare_ndarray ( data ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % type ( data ) ) <EOL> if byteorder is None : <EOL> byteorder = sys . byteorder <EOL> self . _byteorder = _set_endianness ( byteorder ) <EOL> self . _encoding = encoding <EOL> self . _file = get_file_obj ( fname , '<STR_LIT:wb>' , encoding ) <EOL> def _write ( self , to_write ) : <EOL> """<STR_LIT>""" <EOL> self . _file . write ( asbytes ( to_write ) ) <EOL> def _prepare_structured_array ( self , data ) : <EOL> self . nobs = len ( data ) <EOL> self . nvar = len ( data . dtype ) <EOL> self . data = data <EOL> self . datarows = iter ( data ) <EOL> dtype = data . dtype <EOL> descr = dtype . descr <EOL> if dtype . names is None : <EOL> varlist = _default_names ( nvar ) <EOL> else : <EOL> varlist = dtype . names <EOL> convert_dates = self . _convert_dates <EOL> if convert_dates is not None : <EOL> convert_dates = _maybe_convert_to_int_keys ( convert_dates , <EOL> varlist ) <EOL> self . _convert_dates = convert_dates <EOL> for key in convert_dates : <EOL> descr [ key ] = ( <EOL> descr [ key ] [ <NUM_LIT:0> ] , <EOL> _convert_datetime_to_stata_type ( convert_dates [ key ] ) <EOL> ) <EOL> dtype = np . dtype ( descr ) <EOL> self . varlist = varlist <EOL> self . typlist = [ _dtype_to_stata_type ( dtype [ i ] ) <EOL> for i in range ( self . nvar ) ] <EOL> self . fmtlist = [ _dtype_to_default_stata_fmt ( dtype [ i ] ) <EOL> for i in range ( self . nvar ) ] <EOL> if convert_dates is not None : <EOL> for key in convert_dates : <EOL> self . fmtlist [ key ] = convert_dates [ key ] <EOL> def _prepare_ndarray ( self , data ) : <EOL> if data . ndim == <NUM_LIT:1> : <EOL> data = data [ : , None ] <EOL> self . nobs , self . nvar = data . shape <EOL> self . data = data <EOL> self . datarows = iter ( data ) <EOL> dtype = data . dtype <EOL> self . varlist = _default_names ( self . nvar ) <EOL> self . typlist = [ _dtype_to_stata_type ( dtype ) for i in range ( self . nvar ) ] <EOL> self . fmtlist = [ _dtype_to_default_stata_fmt ( dtype ) <EOL> for i in range ( self . nvar ) ] <EOL> def _prepare_pandas ( self , data ) : <EOL> class DataFrameRowIter ( object ) : <EOL> def __init__ ( self , data ) : <EOL> self . data = data <EOL> def __iter__ ( self ) : <EOL> for i , row in data . iterrows ( ) : <EOL> yield row <EOL> data = data . reset_index ( ) <EOL> self . datarows = DataFrameRowIter ( data ) <EOL> self . nobs , self . nvar = data . shape <EOL> self . data = data <EOL> self . varlist = data . columns . tolist ( ) <EOL> dtypes = data . dtypes <EOL> convert_dates = self . _convert_dates <EOL> if convert_dates is not None : <EOL> convert_dates = _maybe_convert_to_int_keys ( convert_dates , <EOL> self . varlist ) <EOL> self . _convert_dates = convert_dates <EOL> for key in convert_dates : <EOL> new_type = _convert_datetime_to_stata_type ( convert_dates [ key ] ) <EOL> dtypes [ key ] = np . dtype ( new_type ) <EOL> self . typlist = [ _dtype_to_stata_type ( dt ) for dt in dtypes ] <EOL> self . fmtlist = [ _dtype_to_default_stata_fmt ( dt ) for dt in dtypes ] <EOL> if convert_dates is not None : <EOL> for key in convert_dates : <EOL> self . fmtlist [ key ] = convert_dates [ key ] <EOL> def write_file ( self ) : <EOL> self . _write_header ( ) <EOL> self . _write_descriptors ( ) <EOL> self . _write_variable_labels ( ) <EOL> self . _write ( _pad_bytes ( "<STR_LIT>" , <NUM_LIT:5> ) ) <EOL> if self . _convert_dates is None : <EOL> self . _write_data_nodates ( ) <EOL> else : <EOL> self . _write_data_dates ( ) <EOL> def _write_header ( self , data_label = None , time_stamp = None ) : <EOL> byteorder = self . _byteorder <EOL> self . _write ( pack ( "<STR_LIT:b>" , <NUM_LIT> ) ) <EOL> self . _write ( byteorder == "<STR_LIT:>>" and "<STR_LIT>" or "<STR_LIT>" ) <EOL> self . _write ( "<STR_LIT>" ) <EOL> self . _write ( "<STR_LIT:\x00>" ) <EOL> self . _write ( pack ( byteorder + "<STR_LIT:h>" , self . nvar ) [ : <NUM_LIT:2> ] ) <EOL> self . _write ( pack ( byteorder + "<STR_LIT:i>" , self . nobs ) [ : <NUM_LIT:4> ] ) <EOL> if data_label is None : <EOL> self . _write ( self . _null_terminate ( _pad_bytes ( "<STR_LIT>" , <NUM_LIT> ) , <EOL> self . _encoding ) ) <EOL> else : <EOL> self . _write ( self . _null_terminate ( _pad_bytes ( data_label [ : <NUM_LIT> ] , <EOL> <NUM_LIT> ) , self . _encoding ) ) <EOL> if time_stamp is None : <EOL> time_stamp = datetime . datetime . now ( ) <EOL> elif not isinstance ( time_stamp , datetime ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . _write ( self . _null_terminate ( <EOL> time_stamp . strftime ( "<STR_LIT>" ) , <EOL> self . _encoding ) ) <EOL> def _write_descriptors ( self , typlist = None , varlist = None , srtlist = None , <EOL> fmtlist = None , lbllist = None ) : <EOL> nvar = self . nvar <EOL> for typ in self . typlist : <EOL> self . _write ( typ ) <EOL> for name in self . varlist : <EOL> name = self . _null_terminate ( name , self . _encoding ) <EOL> name = _pad_bytes ( asstr ( name [ : <NUM_LIT:32> ] ) , <NUM_LIT> ) <EOL> self . _write ( name ) <EOL> srtlist = _pad_bytes ( "<STR_LIT>" , ( <NUM_LIT:2> * ( nvar + <NUM_LIT:1> ) ) ) <EOL> self . _write ( srtlist ) <EOL> for fmt in self . fmtlist : <EOL> self . _write ( _pad_bytes ( fmt , <NUM_LIT> ) ) <EOL> for i in range ( nvar ) : <EOL> self . _write ( _pad_bytes ( "<STR_LIT>" , <NUM_LIT> ) ) <EOL> def _write_variable_labels ( self , labels = None ) : <EOL> nvar = self . nvar <EOL> if labels is None : <EOL> for i in range ( nvar ) : <EOL> self . _write ( _pad_bytes ( "<STR_LIT>" , <NUM_LIT> ) ) <EOL> def _write_data_nodates ( self ) : <EOL> data = self . datarows <EOL> byteorder = self . _byteorder <EOL> TYPE_MAP = self . TYPE_MAP <EOL> typlist = self . typlist <EOL> for row in data : <EOL> for i , var in enumerate ( row ) : <EOL> typ = ord ( typlist [ i ] ) <EOL> if typ <= <NUM_LIT> : <EOL> if len ( var ) < typ : <EOL> var = _pad_bytes ( asstr ( var ) , len ( var ) + <NUM_LIT:1> ) <EOL> self . _write ( var ) <EOL> else : <EOL> try : <EOL> self . _write ( pack ( byteorder + TYPE_MAP [ typ ] , var ) ) <EOL> except struct_error : <EOL> self . _write ( pack ( byteorder + TYPE_MAP [ typ ] , <EOL> _type_converters [ typ ] ( var ) ) ) <EOL> def _write_data_dates ( self ) : <EOL> convert_dates = self . _convert_dates <EOL> data = self . datarows <EOL> byteorder = self . _byteorder <EOL> TYPE_MAP = self . TYPE_MAP <EOL> MISSING_VALUES = self . MISSING_VALUES <EOL> typlist = self . typlist <EOL> for row in data : <EOL> for i , var in enumerate ( row ) : <EOL> typ = ord ( typlist [ i ] ) <EOL> if i in convert_dates : <EOL> var = _datetime_to_stata_elapsed ( var , self . fmtlist [ i ] ) <EOL> if typ <= <NUM_LIT> : <EOL> if isnull ( var ) : <EOL> var = "<STR_LIT>" <EOL> if len ( var ) < typ : <EOL> var = _pad_bytes ( var , len ( var ) + <NUM_LIT:1> ) <EOL> self . _write ( var ) <EOL> else : <EOL> if isnull ( var ) : <EOL> var = MISSING_VALUES [ typ ] <EOL> self . _write ( pack ( byteorder + TYPE_MAP [ typ ] , var ) ) <EOL> def _null_terminate ( self , s , encoding ) : <EOL> null_byte = '<STR_LIT:\x00>' <EOL> if PY3 : <EOL> s += null_byte <EOL> return s . encode ( encoding ) <EOL> else : <EOL> s += null_byte <EOL> return s <EOL> def genfromdta ( fname , missing_flt = - <NUM_LIT> , encoding = None , pandas = False , <EOL> convert_dates = True ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( fname , string_types ) : <EOL> fhd = StataReader ( open ( fname , '<STR_LIT:rb>' ) , missing_values = False , <EOL> encoding = encoding ) <EOL> elif not hasattr ( fname , '<STR_LIT>' ) : <EOL> raise TypeError ( "<STR_LIT>" "<STR_LIT>" % type ( fname ) ) <EOL> else : <EOL> fhd = StataReader ( fname , missing_values = False , encoding = encoding ) <EOL> header = fhd . file_headers ( ) <EOL> types = header [ '<STR_LIT>' ] <EOL> nobs = header [ '<STR_LIT>' ] <EOL> numvars = header [ '<STR_LIT>' ] <EOL> varnames = header [ '<STR_LIT>' ] <EOL> fmtlist = header [ '<STR_LIT>' ] <EOL> dataname = header [ '<STR_LIT>' ] <EOL> labels = header [ '<STR_LIT>' ] <EOL> data = np . zeros ( ( nobs , numvars ) ) <EOL> stata_dta = fhd . dataset ( ) <EOL> dt = np . dtype ( lzip ( varnames , types ) ) <EOL> data = np . zeros ( ( nobs ) , dtype = dt ) <EOL> for rownum , line in enumerate ( stata_dta ) : <EOL> if None in line : <EOL> for i , val in enumerate ( line ) : <EOL> if val is None : <EOL> line [ i ] = missing_flt <EOL> data [ rownum ] = tuple ( line ) <EOL> if pandas : <EOL> from pandas import DataFrame <EOL> data = DataFrame . from_records ( data ) <EOL> if convert_dates : <EOL> cols = np . where ( lmap ( lambda x : x in _date_formats , fmtlist ) ) [ <NUM_LIT:0> ] <EOL> for col in cols : <EOL> i = col <EOL> col = data . columns [ col ] <EOL> data [ col ] = data [ col ] . apply ( _stata_elapsed_date_to_datetime , <EOL> args = ( fmtlist [ i ] , ) ) <EOL> elif convert_dates : <EOL> cols = np . where ( lmap ( lambda x : x in _date_formats , fmtlist ) ) [ <NUM_LIT:0> ] <EOL> dtype = data . dtype . descr <EOL> dtype = [ ( dt [ <NUM_LIT:0> ] , object ) if i in cols else dt for i , dt in <EOL> enumerate ( dtype ) ] <EOL> data = data . astype ( dtype ) <EOL> for col in cols : <EOL> def convert ( x ) : <EOL> return _stata_elapsed_date_to_datetime ( x , fmtlist [ col ] ) <EOL> data [ data . dtype . names [ col ] ] = lmap ( convert , <EOL> data [ data . dtype . names [ col ] ] ) <EOL> return data <EOL> def savetxt ( fname , X , names = None , fmt = '<STR_LIT>' , delimiter = '<STR_LIT:U+0020>' ) : <EOL> """<STR_LIT>""" <EOL> with get_file_obj ( fname , '<STR_LIT:w>' ) as fh : <EOL> X = np . asarray ( X ) <EOL> if X . ndim == <NUM_LIT:1> : <EOL> if X . dtype . names is None : <EOL> X = np . atleast_2d ( X ) . T <EOL> ncol = <NUM_LIT:1> <EOL> else : <EOL> ncol = len ( X . dtype . descr ) <EOL> else : <EOL> ncol = X . shape [ <NUM_LIT:1> ] <EOL> if isinstance ( fmt , ( list , tuple ) ) : <EOL> if len ( fmt ) != ncol : <EOL> raise AttributeError ( '<STR_LIT>' % str ( fmt ) ) <EOL> format = delimiter . join ( fmt ) <EOL> elif isinstance ( fmt , string_types ) : <EOL> if fmt . count ( '<STR_LIT:%>' ) == <NUM_LIT:1> : <EOL> fmt = [ fmt , ] * ncol <EOL> format = delimiter . join ( fmt ) <EOL> elif fmt . count ( '<STR_LIT:%>' ) != ncol : <EOL> raise AttributeError ( '<STR_LIT>' <EOL> % fmt ) <EOL> else : <EOL> format = fmt <EOL> if names is None and X . dtype . names : <EOL> names = X . dtype . names <EOL> if names is not None : <EOL> fh . write ( delimiter . join ( names ) + '<STR_LIT:\n>' ) <EOL> for row in X : <EOL> fh . write ( format % tuple ( row ) + '<STR_LIT:\n>' ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import os <EOL> curdir = os . path . dirname ( os . path . abspath ( __file__ ) ) <EOL> res1 = genfromdta ( curdir + '<STR_LIT>' ) </s>
