diff options
author | Charles Harris <charlesr.harris@gmail.com> | 2018-01-19 20:48:37 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2018-01-19 20:48:37 -0700 |
commit | 9749eec481b97f63520915b044dc5fdd14c71ebd (patch) | |
tree | e8f2449a99d8a2943681aa5326a46c8afd4accb2 | |
parent | d370662e8d21dc72e2145cc2d6d1de1a4ed630f3 (diff) | |
parent | ba3e5856ca20d66aaa675261e82da1814f9cba1f (diff) | |
download | python-numpy-9749eec481b97f63520915b044dc5fdd14c71ebd.tar.gz python-numpy-9749eec481b97f63520915b044dc5fdd14c71ebd.tar.bz2 python-numpy-9749eec481b97f63520915b044dc5fdd14c71ebd.zip |
Merge pull request #10406 from ev-br/print_opts_ctx
ENH: add `np.printoptions`, a context manager
-rw-r--r-- | doc/release/1.15.0-notes.rst | 7 | ||||
-rw-r--r-- | numpy/core/arrayprint.py | 39 | ||||
-rw-r--r-- | numpy/core/tests/test_arrayprint.py | 32 |
3 files changed, 75 insertions, 3 deletions
diff --git a/doc/release/1.15.0-notes.rst b/doc/release/1.15.0-notes.rst index 16a44113c..996cd47b9 100644 --- a/doc/release/1.15.0-notes.rst +++ b/doc/release/1.15.0-notes.rst @@ -15,6 +15,13 @@ New functions * `np.ma.stack`, the `np.stack` array-joining function generalized to masked arrays. +* `np.printoptions`, the context manager which sets print options temporarily + for the scope of the ``with`` block:: + + >>> with np.printoptions(precision=2): + ... print(np.array([2.0])) / 3 + [0.67] + Deprecations ============ diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py index 6af73e6d7..369b3a913 100644 --- a/numpy/core/arrayprint.py +++ b/numpy/core/arrayprint.py @@ -6,8 +6,8 @@ $Id: arrayprint.py,v 1.9 2005/09/13 13:58:44 teoliphant Exp $ from __future__ import division, absolute_import, print_function __all__ = ["array2string", "array_str", "array_repr", "set_string_function", - "set_printoptions", "get_printoptions", "format_float_positional", - "format_float_scientific"] + "set_printoptions", "get_printoptions", "printoptions", + "format_float_positional", "format_float_scientific"] __docformat__ = 'restructuredtext' # @@ -49,7 +49,7 @@ from .numeric import concatenate, asarray, errstate from .numerictypes import (longlong, intc, int_, float_, complex_, bool_, flexible) import warnings - +import contextlib _format_options = { 'edgeitems': 3, # repr N leading and trailing items of each dimension @@ -273,6 +273,39 @@ def get_printoptions(): return _format_options.copy() +@contextlib.contextmanager +def printoptions(*args, **kwargs): + """Context manager for setting print options. + + Set print options for the scope of the `with` block, and restore the old + options at the end. See `set_printoptions` for the full description of + available options. + + Examples + -------- + + >>> with np.printoptions(precision=2): + ... print(np.array([2.0])) / 3 + [0.67] + + The `as`-clause of the `with`-statement gives the current print options: + + >>> with np.printoptions(precision=2) as opts: + ... assert_equal(opts, np.get_printoptions()) + + See Also + -------- + set_printoptions, get_printoptions + + """ + opts = np.get_printoptions() + try: + np.set_printoptions(*args, **kwargs) + yield np.get_printoptions() + finally: + np.set_printoptions(**opts) + + def _leading_trailing(a, edgeitems, index=()): """ Keep only the N-D corners (leading and trailing edges) of an array. diff --git a/numpy/core/tests/test_arrayprint.py b/numpy/core/tests/test_arrayprint.py index 55fb226cc..97b1bd4da 100644 --- a/numpy/core/tests/test_arrayprint.py +++ b/numpy/core/tests/test_arrayprint.py @@ -721,5 +721,37 @@ def test_unicode_object_array(): assert_equal(repr(x), expected) +class TestContextManager(object): + def test_ctx_mgr(self): + # test that context manager actuall works + with np.printoptions(precision=2): + s = str(np.array([2.0]) / 3) + assert_equal(s, '[0.67]') + + def test_ctx_mgr_restores(self): + # test that print options are actually restrored + opts = np.get_printoptions() + with np.printoptions(precision=opts['precision'] - 1, + linewidth=opts['linewidth'] - 4): + pass + assert_equal(np.get_printoptions(), opts) + + def test_ctx_mgr_exceptions(self): + # test that print options are restored even if an exeption is raised + opts = np.get_printoptions() + try: + with np.printoptions(precision=2, linewidth=11): + raise ValueError + except ValueError: + pass + assert_equal(np.get_printoptions(), opts) + + def test_ctx_mgr_as_smth(self): + opts = {"precision": 2} + with np.printoptions(**opts) as ctx: + saved_opts = ctx.copy() + assert_equal({k: saved_opts[k] for k in opts}, opts) + + if __name__ == "__main__": run_module_suite() |