diff options
| -rw-r--r-- | numpy/add_newdocs.py | 48 | ||||
| -rw-r--r-- | numpy/core/numeric.py | 60 | ||||
| -rw-r--r-- | numpy/core/tests/test_numeric.py | 13 |
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() |