ENH: add clear_and_catch_warnings context manager

Add context manager for tests that, on entry, clears record of any
previous warnings from given modules, so that any warning raised in the
context block does get raised, rather than suppressed by a previous
warning causing an entry in the `__warningregistry_`` of the module.

1 files changed, 69 insertions, 1 deletions

diff --git a/numpy/testing/tests/test_utils.py b/numpy/testing/tests/test_utils.py
index 756ea997e..68075fc3d 100644
--- a/numpy/testing/tests/test_utils.py
+++ b/numpy/testing/tests/test_utils.py
@@ -4,7 +4,12 @@ import warnings
 import sys
 
 import numpy as np
-from numpy.testing import *
+from numpy.testing import (
+    assert_equal, assert_array_equal, assert_almost_equal,
+    assert_array_almost_equal, build_err_msg, raises, assert_raises,
+    assert_warns, assert_no_warnings, assert_allclose, assert_approx_equal,
+    assert_array_almost_equal_nulp, assert_array_max_ulp,
+    clear_and_catch_warnings, run_module_suite)
 import unittest
 
 class _GenericTest(object):
@@ -252,6 +257,7 @@ class TestArrayAlmostEqual(_GenericTest, unittest.TestCase):
         assert_array_almost_equal(b, a)
         assert_array_almost_equal(b, b)
 
+
 class TestAlmostEqual(_GenericTest, unittest.TestCase):
     def setUp(self):
         self._assert_func = assert_almost_equal
@@ -688,5 +694,67 @@ class TestULP(unittest.TestCase):
             self.assertRaises(AssertionError,
                                   lambda: assert_array_max_ulp(nan, nzero,
                                                                maxulp=maxulp))
+
+def assert_warn_len_equal(mod, n_in_context):
+    mod_warns = mod.__warningregistry__
+    # Python 3.4 appears to clear any pre-existing warnings of the same type,
+    # when raising warnings inside a catch_warnings block. So, there is a
+    # warning generated by the tests within the context manager, but no
+    # previous warnings.
+    if 'version' in mod_warns:
+        assert_equal(len(mod_warns), 2)  # including 'version'
+    else:
+        assert_equal(len(mod_warns), n_in_context)
+
+
+def _get_fresh_mod():
+    # Get this module, with warning registry empty
+    my_mod = sys.modules[__name__]
+    try:
+        my_mod.__warningregistry__.clear()
+    except AttributeError:
+        pass
+    return my_mod
+
+
+def test_clear_and_catch_warnings():
+    # Initial state of module, no warnings
+    my_mod = _get_fresh_mod()
+    assert_equal(getattr(my_mod, '__warningregistry__', {}), {})
+    with clear_and_catch_warnings(modules=[my_mod]):
+        warnings.simplefilter('ignore')
+        warnings.warn('Some warning')
+    assert_equal(my_mod.__warningregistry__, {})
+    # Without specified modules, don't clear warnings during context
+    with clear_and_catch_warnings():
+        warnings.simplefilter('ignore')
+        warnings.warn('Some warning')
+    assert_warn_len_equal(my_mod, 1)
+    # Confirm that specifying module keeps old warning, does not add new
+    with clear_and_catch_warnings(modules=[my_mod]):
+        warnings.simplefilter('ignore')
+        warnings.warn('Another warning')
+    assert_warn_len_equal(my_mod, 1)
+    # Another warning, no module spec does add to warnings dict, except on
+    # Python 3.4 (see comments in `assert_warn_len_equal`)
+    with clear_and_catch_warnings():
+        warnings.simplefilter('ignore')
+        warnings.warn('Another warning')
+    assert_warn_len_equal(my_mod, 2)
+
+
+class my_cacw(clear_and_catch_warnings):
+    class_modules = (sys.modules[__name__],)
+
+
+def test_clear_and_catch_warnings_inherit():
+    # Test can subclass and add default modules
+    my_mod = _get_fresh_mod()
+    with my_cacw():
+        warnings.simplefilter('ignore')
+        warnings.warn('Some warning')
+    assert_equal(my_mod.__warningregistry__, {})
+
+
 if __name__ == '__main__':
     run_module_suite()