Merge pull request #9330 from nimish/errstate-improvement

ENH: Make errstate a ContextDecorator in Python3

2 files changed, 16 insertions, 2 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 8a8efddf3..1b8f36c3e 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -12,6 +12,7 @@ import operator
 import sys
 import warnings
 import numbers
+import contextlib
 
 import numpy as np
 from . import multiarray
@@ -2990,7 +2991,7 @@ _Unspecified = _unspecified()
 
 
 @set_module('numpy')
-class errstate(object):
+class errstate(contextlib.ContextDecorator):
     """
     errstate(**kwargs)
 
@@ -3000,7 +3001,12 @@ class errstate(object):
     that context to execute with a known error handling behavior. Upon entering
     the context the error handling is set with `seterr` and `seterrcall`, and
     upon exiting it is reset to what it was before.
-
+    
+    ..  versionchanged:: 1.17.0
+        `errstate` is also usable as a function decorator, saving
+        a level of indentation if an entire function is wrapped.
+        See :py:class:`contextlib.ContextDecorator` for more information. 
+    
     Parameters
     ----------
     kwargs : {divide, over, under, invalid}
diff --git a/numpy/core/tests/test_errstate.py b/numpy/core/tests/test_errstate.py
index 670d485c1..0008c4cc8 100644
--- a/numpy/core/tests/test_errstate.py
+++ b/numpy/core/tests/test_errstate.py
@@ -39,3 +39,11 @@ class TestErrstate(object):
             with np.errstate(call=None):
                 assert_(np.geterrcall() is None, 'call is not None')
         assert_(np.geterrcall() is olderrcall, 'call is not olderrcall')
+
+    def test_errstate_decorator(self):
+        @np.errstate(all='ignore')
+        def foo():
+            a = -np.arange(3)
+            a // 0
+            
+        foo()