Add errstate object to be created in new 'with' statement

1 files changed, 9 insertions, 1 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 66a6e3a04..8a5797e8b 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -16,7 +16,7 @@ __all__ = ['newaxis', 'ndarray', 'flatiter', 'ufunc',
            'load', 'loads', 'isscalar', 'binary_repr', 'base_repr',
            'ones', 'identity', 'allclose', 'compare_chararrays', 'putmask',
            'seterr', 'geterr', 'setbufsize', 'getbufsize',
-           'seterrcall', 'geterrcall', 'flatnonzero',
+           'seterrcall', 'geterrcall', 'errstate', 'flatnonzero',
            'Inf', 'inf', 'infty', 'Infinity',
            'nan', 'NaN', 'False_', 'True_', 'bitwise_not',
            'CLIP', 'RAISE', 'WRAP', 'MAXDIMS', 'BUFSIZE', 'ALLOW_THREADS']
@@ -716,6 +716,14 @@ def geterrcall():
     """
     return umath.geterrobj()[2]
 
+class errstate(object):
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+    def __enter__(self):
+        self.oldstate = seterr(**self.kwargs)
+    def __exit__(self, *exc_info):
+        numpy.seterr(**self.oldstate)
+
 def _setdef():
     defval = [UFUNC_BUFSIZE_DEFAULT, ERR_DEFAULT, None]
     umath.seterrobj(defval)