Refactor a bit redirected output print test.

1 files changed, 20 insertions, 31 deletions

diff --git a/numpy/core/tests/test_print.py b/numpy/core/tests/test_print.py
index 23b7b283a..442a5a73c 100644
--- a/numpy/core/tests/test_print.py
+++ b/numpy/core/tests/test_print.py
@@ -72,42 +72,31 @@ def test_complex_types():
         yield check_complex_type, t
 
 # print tests
+def _test_redirected_print(x, tp):
+    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 check_float_type_print(tp):
-    for x in [0, 1,-1, 1e10, 1e20, float('inf'), float('nan'), float('-inf')] :
-        x = float(x)
-        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)
+    for x in [0, 1,-1, 1e10, 1e20, 'inf', 'nan', '-inf'] :
+        _test_redirected_print(float(x), tp)
 
 def check_complex_type_print(tp):
     # We do not create complex with inf/nan directly because the feature is
     # missing in python < 2.6
-    for x in [complex(0), complex(1), complex(-1), complex(1e10), complex(1e20),
-              complex(float('inf'), 1), complex(float('nan'), 1),
-              complex(float('-inf'), 1)] :
-        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)
+    for x in [0, 1, -1, 1e10, 1e20, complex(float('inf'), 1),
+              complex(float('nan'), 1), complex(float('-inf'), 1)] :
+        _test_redirected_print(complex(x), tp)
 
 def test_float_type_print():
     """Check formatting when using print """