<s> """<STR_LIT>""" <EOL> import os <EOL> import warnings <EOL> from datetime import datetime <EOL> from numpy . testing import * <EOL> import numpy as np <EOL> from pandas import DataFrame , isnull <EOL> import pandas . util . testing as ptesting <EOL> from statsmodels . compat . python import BytesIO , asbytes <EOL> import statsmodels . api as sm <EOL> from statsmodels . iolib . foreign import ( StataWriter , genfromdta , <EOL> _datetime_to_stata_elapsed , _stata_elapsed_date_to_datetime ) <EOL> from statsmodels . datasets import macrodata <EOL> from statsmodels . compat . pandas import version as pandas_version <EOL> pandas_old = pandas_version < '<STR_LIT>' <EOL> DECIMAL_4 = <NUM_LIT:4> <EOL> DECIMAL_3 = <NUM_LIT:3> <EOL> curdir = os . path . dirname ( os . path . abspath ( __file__ ) ) <EOL> def test_genfromdta ( ) : <EOL> from . results . macrodata import macrodata_result as res2 <EOL> res1 = genfromdta ( curdir + '<STR_LIT>' ) <EOL> assert_array_equal ( res1 == res2 , True ) <EOL> def test_genfromdta_pandas ( ) : <EOL> from pandas . util . testing import assert_frame_equal <EOL> dta = macrodata . load_pandas ( ) . data <EOL> curdir = os . path . dirname ( os . path . abspath ( __file__ ) ) <EOL> res1 = sm . iolib . genfromdta ( curdir + '<STR_LIT>' , <EOL> pandas = True ) <EOL> res1 = res1 . astype ( float ) <EOL> assert_frame_equal ( res1 , dta ) <EOL> def test_stata_writer_structured ( ) : <EOL> buf = BytesIO ( ) <EOL> dta = macrodata . load ( ) . data <EOL> dtype = dta . dtype <EOL> dta = dta . astype ( np . dtype ( [ ( '<STR_LIT>' , int ) , <EOL> ( '<STR_LIT>' , int ) ] + dtype . descr [ <NUM_LIT:2> : ] ) ) <EOL> writer = StataWriter ( buf , dta ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta2 = genfromdta ( buf ) <EOL> assert_array_equal ( dta , dta2 ) <EOL> def test_stata_writer_array ( ) : <EOL> buf = BytesIO ( ) <EOL> dta = macrodata . load ( ) . data <EOL> dta = DataFrame . from_records ( dta ) <EOL> dta . columns = [ "<STR_LIT>" % i for i in range ( <NUM_LIT:1> , <NUM_LIT:15> ) ] <EOL> writer = StataWriter ( buf , dta . values ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta2 = genfromdta ( buf ) <EOL> dta = dta . to_records ( index = False ) <EOL> assert_array_equal ( dta , dta2 ) <EOL> def test_missing_roundtrip ( ) : <EOL> buf = BytesIO ( ) <EOL> dta = np . array ( [ ( np . nan , np . inf , "<STR_LIT>" ) ] , <EOL> dtype = [ ( "<STR_LIT>" , float ) , ( "<STR_LIT>" , np . float32 ) , <EOL> ( "<STR_LIT>" , "<STR_LIT>" ) ] ) <EOL> writer = StataWriter ( buf , dta ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta = genfromdta ( buf , missing_flt = np . nan ) <EOL> assert_ ( isnull ( dta [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] ) ) <EOL> assert_ ( isnull ( dta [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] ) ) <EOL> assert_ ( dta [ <NUM_LIT:0> ] [ <NUM_LIT:2> ] == asbytes ( "<STR_LIT>" ) ) <EOL> dta = genfromdta ( os . path . join ( curdir , "<STR_LIT>" ) , <EOL> missing_flt = - <NUM_LIT> ) <EOL> assert_ ( np . all ( [ dta [ <NUM_LIT:0> ] [ i ] == - <NUM_LIT> for i in range ( <NUM_LIT:5> ) ] ) ) <EOL> def test_stata_writer_pandas ( ) : <EOL> buf = BytesIO ( ) <EOL> dta = macrodata . load ( ) . data <EOL> dtype = dta . dtype <EOL> dta = dta . astype ( np . dtype ( [ ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) ] + dtype . descr [ <NUM_LIT:2> : ] ) ) <EOL> dta4 = dta . astype ( np . dtype ( [ ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> ( '<STR_LIT>' , '<STR_LIT>' ) ] + dtype . descr [ <NUM_LIT:2> : ] ) ) <EOL> dta = DataFrame . from_records ( dta ) <EOL> dta4 = DataFrame . from_records ( dta4 ) <EOL> writer = StataWriter ( buf , dta ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta2 = genfromdta ( buf ) <EOL> dta5 = DataFrame . from_records ( dta2 ) <EOL> if dta5 . dtypes [ <NUM_LIT:1> ] is np . dtype ( '<STR_LIT>' ) : <EOL> ptesting . assert_frame_equal ( dta . reset_index ( ) , dta5 ) <EOL> else : <EOL> ptesting . assert_frame_equal ( dta4 , dta5 [ dta5 . columns [ <NUM_LIT:1> : ] ] ) <EOL> def test_stata_writer_unicode ( ) : <EOL> pass <EOL> @ dec . skipif ( pandas_old ) <EOL> def test_genfromdta_datetime ( ) : <EOL> results = [ ( datetime ( <NUM_LIT> , <NUM_LIT:11> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT:20> ) , <NUM_LIT> , <EOL> datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:20> ) , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:8> ) , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) , <EOL> datetime ( <NUM_LIT> , <NUM_LIT:7> , <NUM_LIT:1> ) , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) ) , <EOL> ( datetime ( <NUM_LIT> , <NUM_LIT:12> , <NUM_LIT> , <NUM_LIT:20> , <NUM_LIT:3> , <NUM_LIT:20> ) , - <NUM_LIT> , datetime ( <NUM_LIT> , <NUM_LIT:10> , <NUM_LIT:2> ) , <EOL> datetime ( <NUM_LIT> , <NUM_LIT:6> , <NUM_LIT:10> ) , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) , datetime ( <NUM_LIT> , <NUM_LIT:7> , <NUM_LIT:1> ) , <EOL> datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) , datetime ( <NUM_LIT:2> , <NUM_LIT:1> , <NUM_LIT:1> ) ) ] <EOL> with warnings . catch_warnings ( record = True ) as w : <EOL> dta = genfromdta ( os . path . join ( curdir , "<STR_LIT>" ) ) <EOL> assert_ ( len ( w ) == <NUM_LIT:1> ) <EOL> assert_array_equal ( dta [ <NUM_LIT:0> ] . tolist ( ) , results [ <NUM_LIT:0> ] ) <EOL> assert_array_equal ( dta [ <NUM_LIT:1> ] . tolist ( ) , results [ <NUM_LIT:1> ] ) <EOL> with warnings . catch_warnings ( record = True ) : <EOL> dta = genfromdta ( os . path . join ( curdir , "<STR_LIT>" ) , <EOL> pandas = True ) <EOL> assert_array_equal ( dta . iloc [ <NUM_LIT:0> ] . tolist ( ) , results [ <NUM_LIT:0> ] ) <EOL> assert_array_equal ( dta . iloc [ <NUM_LIT:1> ] . tolist ( ) , results [ <NUM_LIT:1> ] ) <EOL> def test_date_converters ( ) : <EOL> ms = [ - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT:100> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> days = [ - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT:50> , <NUM_LIT:0> , <NUM_LIT:50> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> weeks = [ - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , <NUM_LIT:0> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> months = [ - <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT:100> , - <NUM_LIT> , - <NUM_LIT:12> , - <NUM_LIT:11> , <NUM_LIT:0> , <NUM_LIT:11> , <NUM_LIT:12> , <NUM_LIT> , <NUM_LIT:100> , <NUM_LIT> , <NUM_LIT> ] <EOL> quarter = [ - <NUM_LIT:100> , - <NUM_LIT:50> , - <NUM_LIT:5> , - <NUM_LIT:4> , - <NUM_LIT:3> , <NUM_LIT:0> , <NUM_LIT:3> , <NUM_LIT:4> , <NUM_LIT:5> , <NUM_LIT:50> , <NUM_LIT:100> ] <EOL> half = [ - <NUM_LIT:50> , <NUM_LIT> , <NUM_LIT:30> , <NUM_LIT:10> , <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:10> , <NUM_LIT:30> , <NUM_LIT> , <NUM_LIT:50> ] <EOL> year = [ <NUM_LIT:1> , <NUM_LIT:50> , <NUM_LIT> , <NUM_LIT:1000> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> for i in ms : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in days : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in weeks : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in months : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in quarter : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in half : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> for i in year : <EOL> assert_equal ( _datetime_to_stata_elapsed ( <EOL> _stata_elapsed_date_to_datetime ( i , "<STR_LIT>" ) , "<STR_LIT>" ) , i ) <EOL> @ dec . skipif ( pandas_old ) <EOL> def test_datetime_roundtrip ( ) : <EOL> dta = np . array ( [ ( <NUM_LIT:1> , datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) , <NUM_LIT:2> ) , <EOL> ( <NUM_LIT:2> , datetime ( <NUM_LIT> , <NUM_LIT:2> , <NUM_LIT:1> ) , <NUM_LIT:3> ) , <EOL> ( <NUM_LIT:4> , datetime ( <NUM_LIT> , <NUM_LIT:3> , <NUM_LIT:1> ) , <NUM_LIT:5> ) ] , <EOL> dtype = [ ( '<STR_LIT>' , float ) , ( '<STR_LIT>' , object ) , ( '<STR_LIT>' , float ) ] ) <EOL> buf = BytesIO ( ) <EOL> writer = StataWriter ( buf , dta , { "<STR_LIT>" : "<STR_LIT>" } ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta2 = genfromdta ( buf ) <EOL> assert_equal ( dta , dta2 ) <EOL> dta = DataFrame . from_records ( dta ) <EOL> buf = BytesIO ( ) <EOL> writer = StataWriter ( buf , dta , { "<STR_LIT>" : "<STR_LIT>" } ) <EOL> writer . write_file ( ) <EOL> buf . seek ( <NUM_LIT:0> ) <EOL> dta2 = genfromdta ( buf , pandas = True ) <EOL> ptesting . assert_frame_equal ( dta , dta2 . drop ( '<STR_LIT:index>' , axis = <NUM_LIT:1> ) ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import nose <EOL> nose . runmodule ( argv = [ __file__ , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> exit = False ) </s>
