summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorCharles Harris <charlesr.harris@gmail.com>2018-01-19 20:48:37 -0700
committerGitHub <noreply@github.com>2018-01-19 20:48:37 -0700
commit9749eec481b97f63520915b044dc5fdd14c71ebd (patch)
treee8f2449a99d8a2943681aa5326a46c8afd4accb2
parentd370662e8d21dc72e2145cc2d6d1de1a4ed630f3 (diff)
parentba3e5856ca20d66aaa675261e82da1814f9cba1f (diff)
downloadpython-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.rst7
-rw-r--r--numpy/core/arrayprint.py39
-rw-r--r--numpy/core/tests/test_arrayprint.py32
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()