@@ -46,10 +46,25 @@ def test_genfromdta_pandas():
46
46
assert_frame_equal (res1 , dta .astype (float ))
47
47
48
48
49
+ def test_stata_writer_structured ():
50
+ buf = BytesIO ()
51
+ dta = macrodata .load (as_pandas = False ).data
52
+ dtype = dta .dtype
53
+ dt = [('year' , int ), ('quarter' , int )] + dtype .descr [2 :]
54
+ if not PY3 : # Remove unicode
55
+ dt = [(name .encode ('ascii' ), typ ) for name , typ in dt ]
56
+ dta = dta .astype (np .dtype (dt ))
57
+ writer = StataWriter (buf , dta )
58
+ writer .write_file ()
59
+ buf .seek (0 )
60
+ dta2 = genfromdta (buf )
61
+ assert_array_equal (dta , dta2 )
62
+
63
+
49
64
def test_stata_writer_array ():
50
65
buf = BytesIO ()
51
66
dta = macrodata .load (as_pandas = False ).data
52
- dta = DataFrame .from_records (dta , index = 'index' )
67
+ dta = DataFrame .from_records (dta )
53
68
dta .columns = ["v%d" % i for i in range (1 ,15 )]
54
69
writer = StataWriter (buf , dta .values )
55
70
writer .write_file ()
@@ -58,7 +73,6 @@ def test_stata_writer_array():
58
73
dta = dta .to_records (index = False )
59
74
assert_array_equal (dta , dta2 )
60
75
61
-
62
76
def test_missing_roundtrip ():
63
77
buf = BytesIO ()
64
78
dta = np .array ([(np .nan , np .inf , "" )],
0 commit comments