import numpy as np
from numpy.testing import *

import locale
import sys
from StringIO import StringIO

def check_float_type(tp):
    for x in [0, 1,-1, 1e10, 1e20] :
        assert_equal(str(tp(x)), str(float(x)),
                     err_msg='Failed str formatting for type %s' % tp)

def test_float_types():
    """ Check formatting.

        This is only for the str function, and only for simple types.
        The precision of np.float and np.longdouble aren't the same as the
        python float precision.

    """
    for t in [np.float32, np.double, np.longdouble] :
        yield check_float_type, t

def check_nan_inf_float(tp):
    for x in [float('inf'), float('-inf'), float('nan')]:
        assert_equal(str(tp(x)), str(float(x)),
                     err_msg='Failed str formatting for type %s' % tp)

def test_nan_inf_float():
    """ Check formatting.

        This is only for the str function, and only for simple types.
        The precision of np.float and np.longdouble aren't the same as the
        python float precision.

    """
    for t in [np.float32, np.double, np.longdouble] :
        yield check_nan_inf_float, t

def check_complex_type(tp):
    for x in [0, 1,-1, 1e10, 1e20] :
        assert_equal(str(tp(x)), str(complex(x)),
                     err_msg='Failed str formatting for type %s' % tp)
        assert_equal(str(tp(x*1j)), str(complex(x*1j)),
                     err_msg='Failed str formatting for type %s' % tp)
        assert_equal(str(tp(x + x*1j)), str(complex(x + x*1j)),
                     err_msg='Failed str formatting for type %s' % tp)

def test_complex_types():
    """Check formatting.

        This is only for the str function, and only for simple types.
        The precision of np.float and np.longdouble aren't the same as the
        python float precision.

    """
    for t in [np.complex64, np.cdouble, np.clongdouble] :
        yield check_complex_type, t

# print tests
def check_float_type_print(tp):
    for x in [0, 1,-1, 1e10, 1e20, float('inf'), float('nan'), float('-inf')] :
        file = StringIO()
        file_tp = StringIO()
        stdout = sys.stdout
        try:
            sys.stdout = file_tp
            print tp(x)
            sys.stdout = file
            print x
        finally:
            sys.stdout = stdout

        assert_equal(file.getvalue(), file_tp.getvalue(),
                     err_msg='print failed for type%s' % tp)

def test_float_type_print():
    """Check formatting when using print """
    for t in [np.float32, np.double, np.longdouble] :
        yield check_float_type_print, t

# Locale tests: scalar types formatting should be independant of the locale
def has_french_locale():
    curloc = locale.getlocale(locale.LC_NUMERIC)
    try:
        if not sys.platform == 'win32':
            locale.setlocale(locale.LC_NUMERIC, 'fr_FR')
        else:
            locale.setlocale(locale.LC_NUMERIC, 'FRENCH')

        st = True
    except:
        st = False
    finally:
        locale.setlocale(locale.LC_NUMERIC, locale=curloc)

    return st

def _test_locale_independance(tp):
    # XXX: How to query locale on a given system ?

    # French is one language where the decimal is ',' not '.', and should be
    # relatively common on many systems
    curloc = locale.getlocale(locale.LC_NUMERIC)
    try:
        if not sys.platform == 'win32':
            locale.setlocale(locale.LC_NUMERIC, 'fr_FR')
        else:
            locale.setlocale(locale.LC_NUMERIC, 'FRENCH')

        assert_equal(str(tp(1.2)), str(float(1.2)),
                     err_msg='Failed locale test for type %s' % tp)
    finally:
        locale.setlocale(locale.LC_NUMERIC, locale=curloc)

@np.testing.dec.skipif(not has_french_locale(),
                       "Skipping locale test, French locale not found")
def test_locale_single():
    return _test_locale_independance(np.float32)

@np.testing.dec.skipif(not has_french_locale(),
                       "Skipping locale test, French locale not found")
def test_locale_double():
    return _test_locale_independance(np.double)

@np.testing.dec.skipif(not has_french_locale(),
                       "Skipping locale test, French locale not found")
def test_locale_longdouble():
    return _test_locale_independance(np.longdouble)

if __name__ == "__main__":
    run_module_suite()

@@ -2433,157 +2433,1037 @@
 def 
-test_float_type_print():%0A    %22%22%22Check formatting when using print %22%22%22%0A    for t in %5Bnp.float32, np.double, np.longdouble%5D :%0A        yield check_float
+check_complex_type_print(tp):%0A    # We do not create complex with inf/nan directly because the feature is%0A    # missing in python %3C 2.6%0A    for x in %5Bcomplex(0), complex(1), complex(-1), complex(1e10), complex(1e20),%0A              complex(float('inf'), 1), complex(float('nan'), 1),%0A              complex(float('-inf'), 1)%5D :%0A        file = StringIO()%0A        file_tp = StringIO()%0A        stdout = sys.stdout%0A        try:%0A            sys.stdout = file_tp%0A            print tp(x)%0A            sys.stdout = file%0A            print x%0A        finally:%0A            sys.stdout = stdout%0A%0A        assert_equal(file.getvalue(), file_tp.getvalue(),%0A                     err_msg='print failed for type%25s' %25 tp)%0A%0Adef test_float_type_print():%0A    %22%22%22Check formatting when using print %22%22%22%0A    for t in %5Bnp.float32, np.double, np.longdouble%5D :%0A        yield check_float_type_print, t%0A%0Adef test_complex_type_print():%0A    %22%22%22Check formatting when using print %22%22%22%0A    for t in %5Bnp.complex64, np.cdouble, np.clongdouble%5D :%0A        yield check_complex
 _typ

Add print tests for complex types.