Merge pull request #9330 from nimish/errstate-improvement

ENH: Make errstate a ContextDecorator in Python3

3 files changed, 37 insertions, 2 deletions

diff --git a/doc/release/1.17.0-notes.rst b/doc/release/1.17.0-notes.rst
index c79c966c7..73de0b148 100644
--- a/doc/release/1.17.0-notes.rst
+++ b/doc/release/1.17.0-notes.rst
@@ -81,6 +81,27 @@ new function, ``np.ctypeslib.as_ctypes`` now supports a much wider range of
 array types, including structures, booleans, and integers of non-native
 endianness.
 
+`numpy.errstate` is now also function decorator
+-----------------------------------------------
+
+Currently, if you have a function like::
+
+    def foo():
+        pass
+
+and you want to wrap the whole thing in `errstate`, you have to rewrite it like so::
+
+    def foo():
+        with np.errstate(...):
+            pass
+
+but with this change, you can do::
+
+    @np.errstate(...)
+    def foo():
+        pass
+
+thereby saving a level of indentation
 
 Changes
 =======
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()