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

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.
python · segevfiner · Jul 10, 2018 · Dec 1, 2018 · Feb 21, 2022 · Feb 21, 2022
commit 64ceef977edc1c77c02c58d527dc914b15c36807
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 */