Use contextvar for catch_warnings().

nascheme · Dec 21, 2024 · Dec 28, 2024 · Dec 27, 2024 · Dec 27, 2024 · Jan 3, 2025
commit fbd62898085530136ff00d376e44ef7e245fe030
diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -265,6 +265,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(_type_)
         STRUCT_FOR_ID(_uninitialized_submodules)
         STRUCT_FOR_ID(_warn_unawaited_coroutine)
+        STRUCT_FOR_ID(_warnings_context)
         STRUCT_FOR_ID(_xoptions)
         STRUCT_FOR_ID(abs_tol)
         STRUCT_FOR_ID(access)

diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
diff --git a/Lib/test/test_warnings/__init__.py b/Lib/test/test_warnings/__init__.py
@@ -43,15 +43,14 @@ def warnings_state(module):
     except NameError:
         pass
     original_warnings = warning_tests.warnings
-    original_filters = module.filters
+    saved_context, context = module._new_context()
     try:
-        module.filters = original_filters[:]
         module.simplefilter("once")
         warning_tests.warnings = module
         yield
     finally:
         warning_tests.warnings = original_warnings
-        module.filters = original_filters
+        module._set_context(saved_context)
 
 
 class TestWarning(Warning):
@@ -336,15 +335,15 @@ def test_filterwarnings_duplicate_filters(self):
         with original_warnings.catch_warnings(module=self.module):
             self.module.resetwarnings()
             self.module.filterwarnings("error", category=UserWarning)
-            self.assertEqual(len(self.module.filters), 1)
+            self.assertEqual(len(self.module._get_filters()), 1)
             self.module.filterwarnings("ignore", category=UserWarning)
             self.module.filterwarnings("error", category=UserWarning)
             self.assertEqual(
-                len(self.module.filters), 2,
+                len(self.module._get_filters()), 2,
                 "filterwarnings inserted duplicate filter"
             )
             self.assertEqual(
-                self.module.filters[0][0], "error",
+                self.module._get_filters()[0][0], "error",
                 "filterwarnings did not promote filter to "
                 "the beginning of list"
             )
@@ -353,15 +352,15 @@ def test_simplefilter_duplicate_filters(self):
         with original_warnings.catch_warnings(module=self.module):
             self.module.resetwarnings()
             self.module.simplefilter("error", category=UserWarning)
-            self.assertEqual(len(self.module.filters), 1)
+            self.assertEqual(len(self.module._get_filters()), 1)
             self.module.simplefilter("ignore", category=UserWarning)
             self.module.simplefilter("error", category=UserWarning)
             self.assertEqual(
-                len(self.module.filters), 2,
+                len(self.module._get_filters()), 2,
                 "simplefilter inserted duplicate filter"
             )
             self.assertEqual(
-                self.module.filters[0][0], "error",
+                self.module._get_filters()[0][0], "error",
                 "simplefilter did not promote filter to the beginning of list"
             )
 
@@ -373,7 +372,7 @@ def test_append_duplicate(self):
             self.module.simplefilter("error", append=True)
             self.module.simplefilter("ignore", append=True)
             self.module.warn("test_append_duplicate", category=UserWarning)
-            self.assertEqual(len(self.module.filters), 2,
+            self.assertEqual(len(self.module._get_filters()), 2,
                 "simplefilter inserted duplicate filter"
             )
             self.assertEqual(len(w), 0,
@@ -1049,11 +1048,11 @@ def test_issue31416(self):
         # bad warnings.filters or warnings.defaultaction.
         wmod = self.module
         with original_warnings.catch_warnings(module=wmod):
-            wmod.filters = [(None, None, Warning, None, 0)]
+            wmod._get_filters()[:] = [(None, None, Warning, None, 0)]
             with self.assertRaises(TypeError):
                 wmod.warn_explicit('foo', Warning, 'bar', 1)
 
-            wmod.filters = []
+            wmod._get_filters()[:] = []
             with support.swap_attr(wmod, 'defaultaction', None), \
                  self.assertRaises(TypeError):
                 wmod.warn_explicit('foo', Warning, 'bar', 1)
@@ -1191,17 +1190,17 @@ class CatchWarningTests(BaseTest):
 
     def test_catch_warnings_restore(self):
         wmod = self.module
-        orig_filters = wmod.filters
+        orig_filters = wmod._get_filters()
         orig_showwarning = wmod.showwarning
         # Ensure both showwarning and filters are restored when recording
         with wmod.catch_warnings(module=wmod, record=True):
-            wmod.filters = wmod.showwarning = object()
-        self.assertIs(wmod.filters, orig_filters)
+            wmod.get_context()._filters = wmod.showwarning = object()
+        self.assertIs(wmod._get_filters(), orig_filters)
         self.assertIs(wmod.showwarning, orig_showwarning)
         # Same test, but with recording disabled
         with wmod.catch_warnings(module=wmod, record=False):
-            wmod.filters = wmod.showwarning = object()
-        self.assertIs(wmod.filters, orig_filters)
+            wmod.get_context()._filters = wmod.showwarning = object()
+        self.assertIs(wmod._get_filters(), orig_filters)
         self.assertIs(wmod.showwarning, orig_showwarning)
 
     def test_catch_warnings_recording(self):
@@ -1240,21 +1239,21 @@ def test_catch_warnings_reentry_guard(self):
 
     def test_catch_warnings_defaults(self):
         wmod = self.module
-        orig_filters = wmod.filters
+        orig_filters = wmod._get_filters()
         orig_showwarning = wmod.showwarning
         # Ensure default behaviour is not to record warnings
         with wmod.catch_warnings(module=wmod) as w:
             self.assertIsNone(w)
             self.assertIs(wmod.showwarning, orig_showwarning)
-            self.assertIsNot(wmod.filters, orig_filters)
-        self.assertIs(wmod.filters, orig_filters)
+            self.assertIsNot(wmod._get_filters(), orig_filters)
+        self.assertIs(wmod._get_filters(), orig_filters)
         if wmod is sys.modules['warnings']:
             # Ensure the default module is this one
             with wmod.catch_warnings() as w:
                 self.assertIsNone(w)
                 self.assertIs(wmod.showwarning, orig_showwarning)
-                self.assertIsNot(wmod.filters, orig_filters)
-            self.assertIs(wmod.filters, orig_filters)
+                self.assertIsNot(wmod._get_filters(), orig_filters)
+            self.assertIs(wmod._get_filters(), orig_filters)
 
     def test_record_override_showwarning_before(self):
         # Issue #28835: If warnings.showwarning() was overridden, make sure
@@ -1406,7 +1405,7 @@ def test_default_filter_configuration(self):
             code = "import sys; sys.modules.pop('warnings', None); sys.modules['_warnings'] = None; "
         else:
             code = ""
-        code += "import warnings; [print(f) for f in warnings.filters]"
+        code += "import warnings; [print(f) for f in warnings._get_filters()]"
 
         rc, stdout, stderr = assert_python_ok("-c", code, __isolated=True)
         stdout_lines = [line.strip() for line in stdout.splitlines()]
@@ -1821,6 +1820,7 @@ async def coro(self):
         self.assertFalse(inspect.iscoroutinefunction(Cls.sync))
         self.assertTrue(inspect.iscoroutinefunction(Cls.coro))
 
+
 def setUpModule():
     py_warnings.onceregistry.clear()
     c_warnings.onceregistry.clear()

diff --git a/Lib/warnings.py b/Lib/warnings.py
@@ -1,12 +1,73 @@
 """Python part of the warnings subsystem."""
 
 import sys
+import itertools as _itertools
+import contextvars as _contextvars
 
 
 __all__ = ["warn", "warn_explicit", "showwarning",
            "formatwarning", "filterwarnings", "simplefilter",
            "resetwarnings", "catch_warnings", "deprecated"]
 
+class _Context:
+    def __init__(self, filters):
+        self._filters = filters
+        self.log = None  # if set to a list, logging is enabled
+
+    def copy(self):
+        context = _Context(self._filters[:])
+        if self.log is not None:
+            context.log = self.log
+        return context
+
+    def _record_warning(self, msg):
+        self.log.append(msg)
+
+
+class _GlobalContext(_Context):
+    def __init__(self):
+        self.log = None
+
+    @property
+    def _filters(self):
+        # Since there is quite a lot of code that assigns to
+        # warnings.filters, this needs to return the current value of
+        # the module global.
+        try:
+            return filters
+        except NameError:
+            # 'filters' global was deleted.  Do we need to actually handle this case?
+            return []
+
+_global_context = _GlobalContext()
+_warnings_context = _contextvars.ContextVar('warnings_context')
+
+def get_context():
+    try:
+        context = _warnings_context.get()
+    except LookupError:
+        context = _global_context
+        _warnings_context.set(context)
+    return context
+
+
+def _set_context(context):
+    _warnings_context.set(context)
+
+
+def _new_context():
+    old_context = get_context()
+    new_context = old_context.copy()
+    _set_context(new_context)
+    return old_context, new_context
+
+
+def _get_filters():
+    """Return the current list of filters.  This is a non-public API used by
+    the unit tests."""
+    return get_context()._filters
+
+
 def showwarning(message, category, filename, lineno, file=None, line=None):
     """Hook to write a warning to a file; replace if you like."""
     msg = WarningMessage(message, category, filename, lineno, file, line)
@@ -18,6 +79,10 @@ def formatwarning(message, category, filename, lineno, line=None):
     return _formatwarnmsg_impl(msg)
 
 def _showwarnmsg_impl(msg):
+    context = get_context()
+    if context.log is not None:
+        context._record_warning(msg)
+        return
     file = msg.file
     if file is None:
         file = sys.stderr
@@ -129,7 +194,7 @@ def _formatwarnmsg(msg):
     return _formatwarnmsg_impl(msg)
 
 def filterwarnings(action, message="", category=Warning, module="", lineno=0,
-                   append=False):
+                   append=False, *, context=None):
     """Insert an entry into the list of warnings filters (at the front).
 
     'action' -- one of "error", "ignore", "always", "all", "default", "module",
@@ -165,9 +230,11 @@ def filterwarnings(action, message="", category=Warning, module="", lineno=0,
     else:
         module = None
 
-    _add_filter(action, message, category, module, lineno, append=append)
+    _add_filter(action, message, category, module, lineno, append=append,
+                context=context)
 
-def simplefilter(action, category=Warning, lineno=0, append=False):
+def simplefilter(action, category=Warning, lineno=0, append=False, *,
+                 context=None):
     """Insert a simple entry into the list of warnings filters (at the front).
 
     A simple filter matches all modules and messages.
@@ -183,16 +250,20 @@ def simplefilter(action, category=Warning, lineno=0, append=False):
         raise TypeError("lineno must be an int")
     if lineno < 0:
         raise ValueError("lineno must be an int >= 0")
-    _add_filter(action, None, category, None, lineno, append=append)
+    _add_filter(action, None, category, None, lineno, append=append,
+                context=context)
 
 def _filters_mutated():
     # Even though this function is part of the public API, it's used by
     # a fair amount of user code.
     with _lock:
         _filters_mutated_lock_held()
 
-def _add_filter(*item, append):
+def _add_filter(*item, append, context=None):
     with _lock:
+        if context is None:
+            context = get_context()
+        filters = context._filters
         if not append:
             # Remove possible duplicate filters, so new one will be placed
             # in correct place. If append=True and duplicate exists, do nothing.
@@ -206,10 +277,12 @@ def _add_filter(*item, append):
                 filters.append(item)
         _filters_mutated_lock_held()
 
-def resetwarnings():
+def resetwarnings(*, context=None):
     """Clear the list of warning filters, so that no filters are active."""
     with _lock:
-        filters[:] = []
+        if context is None:
+            context = get_context()
+        del context._filters[:]
         _filters_mutated_lock_held()
 
 class _OptionError(Exception):
@@ -378,7 +451,7 @@ def warn_explicit(message, category, filename, lineno,
         if registry.get(key):
             return
         # Search the filters
-        for item in filters:
+        for item in get_context()._filters:
             action, msg, cat, mod, ln = item
             if ((msg is None or msg.match(text)) and
                 issubclass(category, cat) and
@@ -499,31 +572,28 @@ def __enter__(self):
             raise RuntimeError("Cannot enter %r twice" % self)
         self._entered = True
         with _lock:
-            self._filters = self._module.filters
-            self._module.filters = self._filters[:]
-            self._module._filters_mutated_lock_held()
+            self._saved_context, context = self._module._new_context()
             self._showwarning = self._module.showwarning
             self._showwarnmsg_impl = self._module._showwarnmsg_impl
+            if self._record:
+                context.log = log = []
+                # Reset showwarning() to the default implementation to make sure
+                # that _showwarnmsg() calls _showwarnmsg_impl()
+                self._module.showwarning = self._module._showwarning_orig
+            else:
+                log = None
         if self._filter is not None:
-            simplefilter(*self._filter)
-        if self._record:
-            log = []
-            self._module._showwarnmsg_impl = log.append
-            # Reset showwarning() to the default implementation to make sure
-            # that _showwarnmsg() calls _showwarnmsg_impl()
-            self._module.showwarning = self._module._showwarning_orig
-            return log
-        else:
-            return None
+            self._module.simplefilter(*self._filter, context=context)
+        return log
 
     def __exit__(self, *exc_info):
         if not self._entered:
             raise RuntimeError("Cannot exit %r without entering first" % self)
         with _lock:
-            self._module.filters = self._filters
-            self._module._filters_mutated_lock_held()
+            self._module._warnings_context.set(self._saved_context)
             self._module.showwarning = self._showwarning
             self._module._showwarnmsg_impl = self._showwarnmsg_impl
+            self._module._filters_mutated_lock_held()
 
 
 class deprecated:
@@ -762,3 +832,9 @@ def _filters_mutated_lock_held():
         simplefilter("ignore", category=ResourceWarning, append=1)
 
 del _warnings_defaults
+
+#def __getattr__(name):
+#    if name == "filters":
+#        warn('Accessing warnings.filters is likely not thread-safe.', DeprecationWarning, stacklevel=2)
+#        return get_context()._filters
+#    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")