bpo-29672: Save and restore module warning registries in catch_warnings

segevfiner · segevfiner · commit 64ceef977edc · 2018-07-15T19:27:46.000+03:00
This avoids catch_warnings wiping out all module warning registries due
to calling _filters_mutated. Which causes all warnings recorded in those
registries to be shown again.
diff --git a/Lib/test/test_warnings/__init__.py b/Lib/test/test_warnings/__init__.py
@@ -1063,6 +1063,15 @@ def test_check_warnings(self):
             with support.check_warnings(('foo', RuntimeWarning)):
                 wmod.warn("foo")
 
+    def test_check_warnings_restore_registries(self):
+        global __warningregistry__
+        wmod = self.module
+        orig_registry = __warningregistry__ = {}
+        with wmod.catch_warnings(module=wmod):
+            wmod.warn("foo")
+            assert len(__warningregistry__) != 0
+        assert len(__warningregistry__) == 0
+
 class CCatchWarningTests(CatchWarningTests, unittest.TestCase):
     module = c_warnings
 
diff --git a/Lib/warnings.py b/Lib/warnings.py
@@ -116,6 +116,9 @@ def _formatwarnmsg(msg):
                       msg.filename, msg.lineno, line=msg.line)
     return _formatwarnmsg_impl(msg)
 
+def _registrycleared(registry):
+    """Hook that notifies when a module warning registry is cleared."""
+
 def filterwarnings(action, message="", category=Warning, module="", lineno=0,
                    append=False):
     """Insert an entry into the list of warnings filters (at the front).
@@ -328,6 +331,7 @@ def warn_explicit(message, category, filename, lineno,
     if registry is None:
         registry = {}
     if registry.get('version', 0) != _filters_version:
+        _registrycleared(registry)
         registry.clear()
         registry['version'] = _filters_version
     if isinstance(message, Warning):
@@ -438,6 +442,7 @@ def __init__(self, *, record=False, module=None):
         self._record = record
         self._module = sys.modules['warnings'] if module is None else module
         self._entered = False
+        self._old_registries = []
 
     def __repr__(self):
         args = []
@@ -454,7 +459,9 @@ def __enter__(self):
         self._entered = True
         self._filters = self._module.filters
         self._module.filters = self._filters[:]
-        self._module._filters_mutated()
+        self._orig_registrycleared = self._module._registrycleared
+        self._module._registrycleared = self._registrycleared
+        self._filters_version = self._module._filters_mutated()
         self._showwarning = self._module.showwarning
         self._showwarnmsg_impl = self._module._showwarnmsg_impl
         if self._record:
@@ -471,10 +478,17 @@ def __exit__(self, *exc_info):
         if not self._entered:
             raise RuntimeError("Cannot exit %r without entering first" % self)
         self._module.filters = self._filters
-        self._module._filters_mutated()
+        self._module._registrycleared = self._orig_registrycleared
+        self._module._set_filters_version(self._filters_version)
+        for registry, registry_copy in self._old_registries:
+            registry.clear()
+            registry.update(registry_copy)
         self._module.showwarning = self._showwarning
         self._module._showwarnmsg_impl = self._showwarnmsg_impl
 
+    def _registrycleared(self, registry):
+        self._old_registries.append((registry, registry.copy()))
+
 
 # Private utility function called by _PyErr_WarnUnawaitedCoroutine
 def _warn_unawaited_coroutine(coro):
@@ -509,7 +523,8 @@ def extract():
 # If either if the compiled regexs are None, match anything.
 try:
     from _warnings import (filters, _defaultaction, _onceregistry,
-                           warn, warn_explicit, _filters_mutated)
+                           warn, warn_explicit, _filters_mutated,
+                           _set_filters_version)
     defaultaction = _defaultaction
     onceregistry = _onceregistry
     _warnings_defaults = True
@@ -522,7 +537,12 @@ def extract():
 
     def _filters_mutated():
         global _filters_version
+        old_filters_version = _filters_version
         _filters_version += 1
+        return old_filters_version
+
+    def _set_filters_version(filters_version):
+        _filters_version = filters_version
 
     _warnings_defaults = False
 
diff --git a/Python/_warnings.c b/Python/_warnings.c
@@ -243,6 +243,39 @@ get_filter(PyObject *category, PyObject *text, Py_ssize_t lineno,
 }
 
 
+static int
+call_registrycleared(PyObject *registry)
+{
+    PyObject *_registrycleared, *res;
+    _Py_IDENTIFIER(_registrycleared);
+
+    _registrycleared = get_warnings_attr(&PyId__registrycleared, 0);
+    if (_registrycleared == NULL) {
+        if (PyErr_Occurred())
+            return -1;
+        return 0;
+    }
+
+    if (!PyCallable_Check(_registrycleared)) {
+        PyErr_SetString(PyExc_TypeError,
+                "warnings._registrycleared() must be set to a callable");
+        goto error;
+    }
+
+    res = PyObject_CallFunctionObjArgs(_registrycleared, registry, NULL);
+    Py_DECREF(_registrycleared);
+
+    if (res == NULL);
+        return -1;
+
+    Py_DECREF(res);
+    return 0;
+
+error:
+    Py_XDECREF(_registrycleared);
+    return -1;
+}
+
 static int
 already_warned(PyObject *registry, PyObject *key, int should_set)
 {
@@ -256,6 +289,7 @@ already_warned(PyObject *registry, PyObject *key, int should_set)
     if (version_obj == NULL
         || !PyLong_CheckExact(version_obj)
         || PyLong_AsLong(version_obj) != _PyRuntime.warnings.filters_version) {
+        call_registrycleared(registry);
         PyDict_Clear(registry);
         version_obj = PyLong_FromLong(_PyRuntime.warnings.filters_version);
         if (version_obj == NULL)
@@ -910,7 +944,19 @@ warnings_warn_explicit(PyObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 warnings_filters_mutated(PyObject *self, PyObject *args)
 {
-    _PyRuntime.warnings.filters_version++;
+    return PyLong_FromLong(_PyRuntime.warnings.filters_version++);
+}
+
+static PyObject *
+warnings_set_filters_version(PyObject *self, PyObject *args)
+{
+    long filters_version;
+
+    if (!PyArg_ParseTuple(args, "l:_set_filters_version", &filters_version))
+        return NULL;
+
+    _PyRuntime.warnings.filters_version = filters_version;
+
     Py_RETURN_NONE;
 }
 
@@ -1154,6 +1200,8 @@ static PyMethodDef warnings_functions[] = {
         METH_VARARGS | METH_KEYWORDS, warn_explicit_doc},
     {"_filters_mutated", (PyCFunction)warnings_filters_mutated, METH_NOARGS,
         NULL},
+    {"_set_filters_version", (PyCFunction)warnings_set_filters_version,
+        METH_VARARGS, NULL},
     /* XXX(brett.cannon): add showwarning? */
     /* XXX(brett.cannon): Reasonable to add formatwarning? */
     {NULL, NULL}                /* sentinel */
