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

by the warnings.warn(stacklevel=) parameter.

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.

Lib/test/test_warnings/__main__.py

@@ -0,0 +1,3 @@
+import unittest
+
+unittest.main('test.test_warnings')

Lib/test/test_warnings/data/import_warning.py

@@ -0,0 +1,3 @@
+import warnings
+
+warnings.warn('module-level warning', DeprecationWarning, stacklevel=2)

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"

Misc/NEWS

@@ -10,6 +10,9 @@ Release date: 2015-09-06
 Core and Builtins
 -----------------
 
+- Issue #24305: Prevent import subsystem stack frames from being counted
+  by the warnings.warn(stacklevel=) parameter.
+
 - Issue #24912: Prevent __class__ assignment to immutable built-in objects.
 
 - Issue #24975: Fix AST compilation for PEP 448 syntax.

Python/_warnings.c

@@ -513,6 +513,64 @@ warn_explicit(PyObject *category, PyObject *message,
     return result;  /* Py_None or NULL. */
 }
 
+static int
+is_internal_frame(PyFrameObject *frame)
+{
+    static PyObject *importlib_string = NULL;
+    static PyObject *bootstrap_string = NULL;
+    PyObject *filename;
+    int contains;
+
+    if (importlib_string == NULL) {
+        importlib_string = PyUnicode_FromString("importlib");
+        if (importlib_string == NULL) {
+            return 0;
+        }
+
+        bootstrap_string = PyUnicode_FromString("_bootstrap");
+        if (bootstrap_string == NULL) {
+            Py_DECREF(importlib_string);
+            return 0;
+        }
+        Py_INCREF(importlib_string);
+        Py_INCREF(bootstrap_string);
+    }
+
+    if (frame == NULL || frame->f_code == NULL ||
+            frame->f_code->co_filename == NULL) {
+        return 0;
+    }
+    filename = frame->f_code->co_filename;
+    if (!PyUnicode_Check(filename)) {
+        return 0;
+    }
+    contains = PyUnicode_Contains(filename, importlib_string);
+    if (contains < 0) {
+        return 0;
+    }
+    else if (contains > 0) {
+        contains = PyUnicode_Contains(filename, bootstrap_string);
+        if (contains < 0) {
+            return 0;
+        }
+        else if (contains > 0) {
+            return 1;
+        }
+    }
+
+    return 0;
+}
+
+static PyFrameObject *
+next_external_frame(PyFrameObject *frame)
+{
+    do {
+        frame = frame->f_back;
+    } while (frame != NULL && is_internal_frame(frame));
+
+    return frame;
+}
+
 /* filename, module, and registry are new refs, globals is borrowed */
 /* Returns 0 on error (no new refs), 1 on success */
 static int
@@ -523,8 +581,18 @@ setup_context(Py_ssize_t stack_level, PyObject **filename, int *lineno,
 
     /* Setup globals and lineno. */
     PyFrameObject *f = PyThreadState_GET()->frame;
-    while (--stack_level > 0 && f != NULL)
-        f = f->f_back;
+    // Stack level comparisons to Python code is off by one as there is no
+    // warnings-related stack level to avoid.
+    if (stack_level <= 0 || is_internal_frame(f)) {
+        while (--stack_level > 0 && f != NULL) {
+            f = f->f_back;
+        }
+    }
+    else {
+        while (--stack_level > 0 && f != NULL) {
+            f = next_external_frame(f);
+        }
+    }
 
     if (f == NULL) {
         globals = PyThreadState_Get()->interp->sysdict;