From 3c8ac112455ab9d1792101ac25b8fa1472fac28b Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Mon, 15 Jan 2018 03:00:21 +0300
Subject: ENH: add `np.printoptions`, a context manager

Just a syntax sugar over the get_printoptions/set_printoptions pair.
---
 numpy/core/arrayprint.py | 40 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

(limited to 'numpy/core/arrayprint.py')

diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py
index 6af73e6d7..d6e3cfda3 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,40 @@ def get_printoptions():
     return _format_options.copy()
 
 
+@contextlib.contextmanager
+def printoptions(*args, **kwds):
+    """Context manager for setting print options.
+
+    See `set_printoptions` for the full description of available options.
+
+    Examples
+    --------
+
+    `printoptions` sets print options temporarily, and restores them back
+    at the end of the `with`-block:
+
+    >>> 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, **kwds)
+        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.
-- 
cgit v1.2.1