"""
Tests for iolib/foreign.py
"""

from numpy.testing import *
import numpy as np
import statsmodels.api as sm
import os
from statsmodels.iolib.foreign import StataWriter, genfromdta
from statsmodels.datasets import macrodata

# Test precisions
DECIMAL_4 = 4
DECIMAL_3 = 3


def test_genfromdta():
    #Test genfromdta vs. results/macrodta.npy created with genfromtxt.
    #NOTE: Stata handles data very oddly.  Round tripping from csv to dta
    #    to ndarray 2710.349 (csv) -> 2510.2491 (stata) -> 2710.34912109375
    #    (dta/ndarray)
    curdir = os.path.dirname(os.path.abspath(__file__))
    #res2 = np.load(curdir+'/results/macrodata.npy')
    #res2 = res2.view((float,len(res2[0])))
    from results.macrodata import macrodata_result as res2
    res1 = genfromdta(curdir+'/../../datasets/macrodata/macrodata.dta')
    #res1 = res1.view((float,len(res1[0])))
    assert_array_equal(res1 == res2, True)

def test_genfromdta_pandas():
    from pandas.util.testing import assert_frame_equal
    dta = macrodata.load_pandas().data
    curdir = os.path.dirname(os.path.abspath(__file__))
    res1 = sm.iolib.genfromdta(curdir+'/../../datasets/macrodata/macrodata.dta',
                        pandas=True)
    res1 = res1.astype(float)
    assert_frame_equal(res1, dta)

def test_stata_writer():
    from StringIO import StringIO
    buf = StringIO()
    dta = macrodata.load().data
    dtype = dta.dtype
    dta = dta.astype(np.dtype([('year', 'i8'),
                               ('quarter', 'i4')] + dtype.descr[2:]))
    writer = StataWriter(buf, dta)
    writer.write_file()
    buf.seek(0)
    dta2 = genfromdta(buf)
    assert_array_equal(dta, dta2)


if __name__ == "__main__":
    import nose
    nose.runmodule(argv=[__file__,'-vvs','-x','--pdb'],
                       exit=False)

@@ -29,16 +29,46 @@
 py%0A%22%22%22%0A%0A
+from StringIO import StringIO%0A
 from num
@@ -87,16 +87,16 @@
 mport *%0A
-
 import n
@@ -249,16 +249,84 @@
 acrodata
+%0Afrom pandas import DataFrame%0Aimport pandas.util.testing as ptesting
 %0A%0A# Test
@@ -1394,45 +1394,22 @@
 iter
-():%0A    from StringIO import StringIO
+_structured():
 %0A   
@@ -1699,16 +1699,16 @@
 ta(buf)%0A
-
     asse
@@ -1734,16 +1734,1065 @@
  dta2)%0A%0A
+def test_stata_writer_array():%0A    buf = StringIO()%0A    dta = macrodata.load().data%0A    dta = DataFrame.from_records(dta)%0A    dta.columns = %5B%22v%25d%22 %25 i for i in range(1,15)%5D%0A    writer = StataWriter(buf, dta.values)%0A    writer.write_file()%0A    buf.seek(0)%0A    dta2 = genfromdta(buf)%0A    dta = dta.to_records(index=False)%0A    assert_array_equal(dta, dta2)%0A%0Adef test_stata_writer_missing():%0A    pass%0A%0Adef test_stata_writer_pandas():%0A    buf = StringIO()%0A    dta = macrodata.load().data%0A    dtype = dta.dtype%0A    #as of 0.9.0 pandas only supports i8 and f8%0A    dta = dta.astype(np.dtype(%5B('year', 'i8'),%0A                               ('quarter', 'i8')%5D + dtype.descr%5B2:%5D))%0A    dta = DataFrame.from_records(dta)%0A    writer = StataWriter(buf, dta)%0A    writer.write_file()%0A    buf.seek(0)%0A    dta2 = genfromdta(buf)%0A    ptesting.assert_frame_equal(dta.reset_index(), DataFrame.from_records(dta2))%0A%0A%0Adef test_stata_writer_unicode():%0A    # make sure to test with characters outside the latin-1 encoding%0A    pass%0A%0Adef test_stata_writer_datetime():%0A    pass%0A%0A
 %0Aif __na

Add some StataWriter tests