BUG: core: make set_string_function(None) to restore the functions to what they are on import (#1130)

3 files changed, 73 insertions, 48 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 350461208..e7899b1e7 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -923,53 +923,7 @@ add_newdoc('numpy.core.multiarray', 'set_string_function',
     """
     set_string_function(f, repr=1)
 
-    Set a Python function to be used when pretty printing arrays.
-
-    Parameters
-    ----------
-    f : function or None
-        Function to be used to pretty print arrays. The function should expect
-        a single array argument and return a string of the representation of
-        the array. If None, the function is reset to the default NumPy function
-        to print arrays.
-    repr : bool, optional
-        If True (default), the function for pretty printing (``__repr__``)
-        is set, if False the function that returns the default string
-        representation (``__str__``) is set.
-
-    See Also
-    --------
-    set_printoptions, get_printoptions
-
-    Examples
-    --------
-    >>> def pprint(arr):
-    ...     return 'HA! - What are you going to do now?'
-    ...
-    >>> np.set_string_function(pprint)
-    >>> a = np.arange(10)
-    >>> a
-    HA! - What are you going to do now?
-    >>> print a
-    [0 1 2 3 4 5 6 7 8 9]
-
-    We can reset the function to the default:
-
-    >>> np.set_string_function(None)
-    >>> a
-    array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 'l')
-
-    `repr` affects either pretty printing or normal string representation.
-    Note that ``__repr__`` is still affected by setting ``__str__``
-    because the width of each array element in the returned string becomes
-    equal to the length of the result of ``__str__()``.
-
-    >>> x = np.arange(4)
-    >>> np.set_string_function(lambda x:'random', repr=False)
-    >>> x.__str__()
-    'random'
-    >>> x.__repr__()
-    'array([     0,      1,      2,      3])'
+    Internal method to set a function to be used when pretty printing arrays.
 
     """)
 
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index eb028874a..7df88e875 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -1430,7 +1430,65 @@ def array_str(a, max_line_width=None, precision=None, suppress_small=None):
     """
     return array2string(a, max_line_width, precision, suppress_small, ' ', "", str)
 
-set_string_function = multiarray.set_string_function
+def set_string_function(f, repr=True):
+    """
+    Set a Python function to be used when pretty printing arrays.
+
+    Parameters
+    ----------
+    f : function or None
+        Function to be used to pretty print arrays. The function should expect
+        a single array argument and return a string of the representation of
+        the array. If None, the function is reset to the default NumPy function
+        to print arrays.
+    repr : bool, optional
+        If True (default), the function for pretty printing (``__repr__``)
+        is set, if False the function that returns the default string
+        representation (``__str__``) is set.
+
+    See Also
+    --------
+    set_printoptions, get_printoptions
+
+    Examples
+    --------
+    >>> def pprint(arr):
+    ...     return 'HA! - What are you going to do now?'
+    ...
+    >>> np.set_string_function(pprint)
+    >>> a = np.arange(10)
+    >>> a
+    HA! - What are you going to do now?
+    >>> print a
+    [0 1 2 3 4 5 6 7 8 9]
+
+    We can reset the function to the default:
+
+    >>> np.set_string_function(None)
+    >>> a
+    array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+
+    `repr` affects either pretty printing or normal string representation.
+    Note that ``__repr__`` is still affected by setting ``__str__``
+    because the width of each array element in the returned string becomes
+    equal to the length of the result of ``__str__()``.
+
+    >>> x = np.arange(4)
+    >>> np.set_string_function(lambda x:'random', repr=False)
+    >>> x.__str__()
+    'random'
+    >>> x.__repr__()
+    'array([     0,      1,      2,      3])'
+
+    """
+    if f is None:
+        if repr:
+            return multiarray.set_string_function(array_repr, 1)
+        else:
+            return multiarray.set_string_function(array_str, 0)
+    else:
+        return multiarray.set_string_function(f, repr)
+
 set_string_function(array_str, 0)
 set_string_function(array_repr, 1)
 
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index 59c195010..8bb6c7315 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -960,5 +960,18 @@ class TestArgwhere:
     def test_list(self):
         assert_equal(np.argwhere([4, 0, 2, 1, 3]), [[0], [2], [3], [4]])
 
+class TestStringFunction:
+    def test_set_string_function(self):
+        a = np.array([1])
+        np.set_string_function(lambda x: "FOO", repr=True)
+        assert_equal(repr(a), "FOO")
+        np.set_string_function(None, repr=True)
+        assert_equal(repr(a), "array([1])")
+
+        np.set_string_function(lambda x: "FOO", repr=False)
+        assert_equal(str(a), "FOO")
+        np.set_string_function(None, repr=False)
+        assert_equal(str(a), "[1]")
+
 if __name__ == "__main__":
     run_module_suite()