Issue #24305: Prevent import subsystem stack frames from being counted

by the warnings.warn(stacklevel=) parameter.

5 files changed, 54 insertions, 14 deletions

diff --git a/Lib/test/test_warnings.py b/Lib/test/test_warnings/__init__.py
index c7d2e5c..991a249 100644
--- a/Lib/test/test_warnings.py
+++ b/Lib/test/test_warnings/__init__.py
@@ -7,7 +7,7 @@ import unittest
 from test import support
 from test.support.script_helper import assert_python_ok, assert_python_failure
 
-from test import warning_tests
+from test.test_warnings.data import stacklevel as warning_tests
 
 import warnings as original_warnings
 
@@ -188,11 +188,11 @@ class FilterTests(BaseTest):
             self.module.resetwarnings()
             self.module.filterwarnings("once", category=UserWarning)
             message = UserWarning("FilterTests.test_once")
-            self.module.warn_explicit(message, UserWarning, "test_warnings.py",
+            self.module.warn_explicit(message, UserWarning, "__init__.py",
                                     42)
             self.assertEqual(w[-1].message, message)
             del w[:]
-            self.module.warn_explicit(message, UserWarning, "test_warnings.py",
+            self.module.warn_explicit(message, UserWarning, "__init__.py",
                                     13)
             self.assertEqual(len(w), 0)
             self.module.warn_explicit(message, UserWarning, "test_warnings2.py",
@@ -298,10 +298,10 @@ class WarnTests(BaseTest):
                     module=self.module) as w:
                 warning_tests.inner("spam1")
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "warning_tests.py")
+                                    "stacklevel.py")
                 warning_tests.outer("spam2")
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "warning_tests.py")
+                                    "stacklevel.py")
 
     def test_stacklevel(self):
         # Test stacklevel argument
@@ -311,25 +311,36 @@ class WarnTests(BaseTest):
                     module=self.module) as w:
                 warning_tests.inner("spam3", stacklevel=1)
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "warning_tests.py")
+                                    "stacklevel.py")
                 warning_tests.outer("spam4", stacklevel=1)
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "warning_tests.py")
+                                    "stacklevel.py")
 
                 warning_tests.inner("spam5", stacklevel=2)
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "test_warnings.py")
+                                    "__init__.py")
                 warning_tests.outer("spam6", stacklevel=2)
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "warning_tests.py")
+                                    "stacklevel.py")
                 warning_tests.outer("spam6.5", stacklevel=3)
                 self.assertEqual(os.path.basename(w[-1].filename),
-                                    "test_warnings.py")
+                                    "__init__.py")
 
                 warning_tests.inner("spam7", stacklevel=9999)
                 self.assertEqual(os.path.basename(w[-1].filename),
                                     "sys")
 
+    def test_stacklevel_import(self):
+        # Issue #24305: With stacklevel=2, module-level warnings should work.
+        support.unload('test.test_warnings.data.import_warning')
+        with warnings_state(self.module):
+            with original_warnings.catch_warnings(record=True,
+                    module=self.module) as w:
+                self.module.simplefilter('always')
+                import test.test_warnings.data.import_warning
+                self.assertEqual(len(w), 1)
+                self.assertEqual(w[0].filename, __file__)
+
     def test_missing_filename_not_main(self):
         # If __file__ is not specified and __main__ is not the module name,
         # then __file__ should be set to the module name.
diff --git a/Lib/test/test_warnings/__main__.py b/Lib/test/test_warnings/__main__.py
new file mode 100644
index 0000000..44e52ec
--- /dev/null
+++ b/Lib/test/test_warnings/__main__.py
@@ -0,0 +1,3 @@
+import unittest
+
+unittest.main('test.test_warnings')
diff --git a/Lib/test/test_warnings/data/import_warning.py b/Lib/test/test_warnings/data/import_warning.py
new file mode 100644
index 0000000..d6ea2ce
--- /dev/null
+++ b/Lib/test/test_warnings/data/import_warning.py
@@ -0,0 +1,3 @@
+import warnings
+
+warnings.warn('module-level warning', DeprecationWarning, stacklevel=2)
\ No newline at end of file
diff --git a/Lib/test/warning_tests.py b/Lib/test/test_warnings/data/stacklevel.py
index d0519ef..d0519ef 100644
--- a/Lib/test/warning_tests.py
+++ b/Lib/test/test_warnings/data/stacklevel.py
diff --git a/Lib/warnings.py b/Lib/warnings.py
index 16246b4..1d4fb20 100644
--- a/Lib/warnings.py
+++ b/Lib/warnings.py
@@ -160,6 +160,20 @@ def _getcategory(category):
     return cat
 
 
+def _is_internal_frame(frame):
+    """Signal whether the frame is an internal CPython implementation detail."""
+    filename = frame.f_code.co_filename
+    return 'importlib' in filename and '_bootstrap' in filename
+
+
+def _next_external_frame(frame):
+    """Find the next frame that doesn't involve CPython internals."""
+    frame = frame.f_back
+    while frame is not None and _is_internal_frame(frame):
+        frame = frame.f_back
+    return frame
+
+
 # Code typically replaced by _warnings
 def warn(message, category=None, stacklevel=1):
     """Issue a warning, or maybe ignore it or raise an exception."""
@@ -174,13 +188,23 @@ def warn(message, category=None, stacklevel=1):
                         "not '{:s}'".format(type(category).__name__))
     # Get context information
     try:
-        caller = sys._getframe(stacklevel)
+        if stacklevel <= 1 or _is_internal_frame(sys._getframe(1)):
+            # If frame is too small to care or if the warning originated in
+            # internal code, then do not try to hide any frames.
+            frame = sys._getframe(stacklevel)
+        else:
+            frame = sys._getframe(1)
+            # Look for one frame less since the above line starts us off.
+            for x in range(stacklevel-1):
+                frame = _next_external_frame(frame)
+                if frame is None:
+                    raise ValueError
     except ValueError:
         globals = sys.__dict__
         lineno = 1
     else:
-        globals = caller.f_globals
-        lineno = caller.f_lineno
+        globals = frame.f_globals
+        lineno = frame.f_lineno
     if '__name__' in globals:
         module = globals['__name__']
     else:
@@ -374,7 +398,6 @@ try:
     defaultaction = _defaultaction
     onceregistry = _onceregistry
     _warnings_defaults = True
-
 except ImportError:
     filters = []
     defaultaction = "default"