Use defining class names for method frames on Python 3.11+

pyinstrument/low_level/stat_profile.c

@@ -190,6 +190,9 @@ static PyObject *whatstrings[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NUL
 static PyObject *SELF_STRING = NULL;
 static PyObject *CLS_STRING = NULL;
 static PyObject *TRACEBACKHIDE_STRING = NULL;
+#if PY_VERSION_HEX >= 0x030b0000
+static PyObject *LOCALS_SUFFIX_STRING = NULL;
+#endif
 
 #define TIMER_TYPE_WALLTIME 0
 #define TIMER_TYPE_WALLTIME_THREAD 1
@@ -228,6 +231,11 @@ stat_profile_init(void)
     CLS_STRING = PyUnicode_InternFromString("cls");
     if (CLS_STRING == NULL) return -1;
 
+#if PY_VERSION_HEX >= 0x030b0000
+    LOCALS_SUFFIX_STRING = PyUnicode_InternFromString(".<locals>");
+    if (LOCALS_SUFFIX_STRING == NULL) return -1;
+#endif
+
     TRACEBACKHIDE_STRING = PyUnicode_InternFromString("__tracebackhide__");
     if (TRACEBACKHIDE_STRING == NULL) return -1;
 
@@ -287,107 +295,54 @@ local_names_from_code(PyCodeObject *code)
 #endif
 }
 
-/**
- * Returns a new reference to the qualified name of a type.
- */
-static PyObject *
-_get_type_qualname(PyTypeObject *type) {
-#if PY_VERSION_HEX >= 0x030b0000 // Python 3.11.0
-    PyObject *qualname = PyType_GetQualName(type);
-#else
-    PyObject *qualname = PyObject_GetAttrString((PyObject *)type, "__qualname__");
-#endif
-    if (qualname == NULL) {
-        return NULL;
-    }
-    if (!PyUnicode_Check(qualname)) {
-        Py_DECREF(qualname);
-        PyErr_SetString(PyExc_TypeError, "type __qualname__ must be a string");
-        return NULL;
-    }
-    return qualname;
-}
-
 #if PY_VERSION_HEX >= 0x030b0000 // Python 3.11.0
 /**
- * Returns a new reference to the qualified name of the class in the frame,
+ * Returns a new reference to the statically defined class name in the frame,
  * or NULL when the frame does not represent a method.
  */
 static PyObject *
 _get_class_name_of_frame(PyFrameObject *frame, PyCodeObject *code) {
-    PyObject *localsNames = PyCode_GetVarnames(code);
-
-    if (localsNames == NULL) {
+    // Class bodies also have qualified names, but are not method calls.
+    if (!(code->co_flags & CO_NEWLOCALS)) {
         return NULL;
     }
 
-    if (PyTuple_GET_SIZE(localsNames) == 0) {
-        Py_DECREF(localsNames);
+    PyObject *qualname = code->co_qualname;
+    Py_ssize_t last_dot = PyUnicode_FindChar(
+        qualname, '.', 0, PyUnicode_GET_LENGTH(qualname), -1
+    );
+    if (last_dot <= 0) {
         return NULL;
     }
 
-    PyObject *firstArgName = PyTuple_GET_ITEM(localsNames, 0);
-
-    if (firstArgName == NULL) {
-        Py_DECREF(localsNames);
+    // A nested function's enclosing scope is not a class.
+    int is_local = PyUnicode_Tailmatch(qualname, LOCALS_SUFFIX_STRING, 0, last_dot, 1);
+    if (is_local != 0) {
         return NULL;
     }
 
-    int has_self = PyUnicode_Compare(firstArgName, SELF_STRING) == 0;
-    int has_cls = PyUnicode_Compare(firstArgName, CLS_STRING) == 0;
+    return PyUnicode_Substring(qualname, 0, last_dot);
+}
 
-    Py_DECREF(localsNames);
+#else
 
-    if (!has_self && !has_cls) {
-        // PyFrame_GetLocals is expensive and changes the frame, so we don't
-        // want to call it unless we have to.
+/**
+ * Returns a new reference to the qualified name of a type.
+ */
+static PyObject *
+_get_type_qualname(PyTypeObject *type) {
+    PyObject *qualname = PyObject_GetAttrString((PyObject *)type, "__qualname__");
+    if (qualname == NULL) {
         return NULL;
     }
-
-    PyObject *result = NULL;
-
-    PyObject *locals = PyFrame_GetLocals(frame);
-
-    if (locals == NULL || !PyMapping_Check(locals)) {
-        Py_XDECREF(locals);
+    if (!PyUnicode_Check(qualname)) {
+        Py_DECREF(qualname);
+        PyErr_SetString(PyExc_TypeError, "type __qualname__ must be a string");
         return NULL;
     }
-
-    // we still have to check the locals has the key, because it could have
-    // been "del'd"
-    if (has_self && PyMapping_HasKey(locals, SELF_STRING)) {
-        PyObject *self = PyObject_GetItem(locals, SELF_STRING);
-
-        if (!self) {
-            PyErr_Clear();
-            Py_DECREF(locals);
-            return NULL;
-        }
-
-        result = _get_type_qualname(Py_TYPE(self));
-        Py_DECREF(self);
-    }
-    else if (has_cls && PyMapping_HasKey(locals, CLS_STRING)) {
-        PyObject *cls = PyObject_GetItem(locals, CLS_STRING);
-
-        if (!cls) {
-            PyErr_Clear();
-            Py_DECREF(locals);
-            return NULL;
-        }
-
-        if (PyType_Check(cls)) {
-            result = _get_type_qualname((PyTypeObject *)cls);
-        }
-        Py_DECREF(cls);
-    }
-
-    Py_DECREF(locals);
-    return result;
+    return qualname;
 }
 
-#else
-
 static PyObject *
 _get_first_arg_from_cell_variables(PyFrameObject *frame, PyCodeObject *code) {
     if (!code->co_cell2arg) {

pyinstrument/low_level/stat_profile_python.py

@@ -4,6 +4,7 @@
 import sys
 import timeit
 import types
+from inspect import CO_NEWLOCALS
 from typing import Any, Callable, List, Optional, Type
 
 from pyinstrument.low_level.pyi_timing_thread_python import (
@@ -114,22 +115,30 @@ def setstatprofile(
 
 
 def get_frame_info(frame: types.FrameType) -> str:
+    code = frame.f_code
     frame_info = "%s\x00%s\x00%i" % (
-        frame.f_code.co_name,
-        frame.f_code.co_filename,
-        frame.f_code.co_firstlineno,
+        code.co_name,
+        code.co_filename,
+        code.co_firstlineno,
     )
 
     class_name = None
-    # try to find self argument for methods
-    self = frame.f_locals.get("self", None)
-    if self and hasattr(self, "__class__") and hasattr(self.__class__, "__qualname__"):
-        class_name = self.__class__.__qualname__
+    if sys.version_info >= (3, 11):
+        if code.co_flags & CO_NEWLOCALS and "." in code.co_qualname:
+            prefix, _, _ = code.co_qualname.rpartition(".")
+            # A nested function's enclosing scope is not a class.
+            if prefix and not prefix.endswith(".<locals>"):
+                class_name = prefix
     else:
-        # also try to find cls argument for class methods
-        cls = frame.f_locals.get("cls", None)
-        if cls and hasattr(cls, "__qualname__"):
-            class_name = cls.__qualname__
+        # try to find self argument for methods
+        self = frame.f_locals.get("self", None)
+        if self and hasattr(self, "__class__") and hasattr(self.__class__, "__qualname__"):
+            class_name = self.__class__.__qualname__
+        else:
+            # also try to find cls argument for class methods
+            cls = frame.f_locals.get("cls", None)
+            if cls and hasattr(cls, "__qualname__"):
+                class_name = cls.__qualname__
 
     frame_hidden = "__tracebackhide__" in frame.f_locals
 

test/low_level/test_frame_info.py

@@ -1,4 +1,5 @@
 import inspect
+import sys
 
 import pytest
 
@@ -113,3 +114,101 @@ def test_frame_info_with_classes(test_function, del_local):
     py_frame_info = test_function(stat_profile_python.get_frame_info, del_local=del_local)
 
     assert c_frame_info == py_frame_info
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+@pytest.mark.parametrize("del_local", [True, False])
+@pytest.mark.parametrize(
+    "method",
+    [
+        nested_instance.get_frame_info_for_a_method,
+        OuterClass.InnerClass.get_frame_info_for_a_class_method,
+        nested_instance.get_frame_info_with_cell_variable,
+        OuterClass.InnerClass.get_frame_info_for_a_class_method_where_cls_is_reassigned,
+    ],
+)
+def test_inherited_method_uses_defining_class(getter, del_local, method):
+    info = method(getter, del_local)
+    assert "\x01cAClass\x01" in info
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+def test_qualified_class_names(getter):
+    class Outer:
+        class Inner甲:
+            @staticmethod
+            def method():
+                return inspect.currentframe()
+
+    class Child(Outer.Inner甲):
+        pass
+
+    frame = Child.method()
+    assert frame
+    assert "\x01ctest_qualified_class_names.<locals>.Outer.Inner甲\x01" in getter(frame)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+def test_nested_function_has_no_class_name(getter):
+    def inner(self):
+        return inspect.currentframe()
+
+    frame = inner(instance)
+    assert frame
+    assert "\x01c" not in getter(frame)
+
+
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+def test_class_body_has_no_class_name(getter):
+    class Outer:
+        class Inner:
+            frame = inspect.currentframe()
+
+    assert Outer.Inner.frame
+    assert "\x01c" not in getter(Outer.Inner.frame)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+def test_metaclass_method_uses_defining_class(getter):
+    class Meta(type):
+        def method(cls):
+            return inspect.currentframe()
+
+    class Target(metaclass=Meta):
+        pass
+
+    frame = Target.method()
+    assert frame
+    assert "\x01ctest_metaclass_method_uses_defining_class.<locals>.Meta\x01" in getter(frame)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+@pytest.mark.parametrize(
+    "getter", [stat_profile_c.get_frame_info, stat_profile_python.get_frame_info]
+)
+def test_function_inside_method_has_no_class_name(getter):
+    class Example:
+        def method(self):
+            def inner():
+                if self:
+                    return inspect.currentframe()
+
+            return inner()
+
+    frame = Example().method()
+    assert frame
+    assert "\x01c" not in getter(frame)

test/test_profiler.py

@@ -2,6 +2,7 @@
 import dataclasses
 import inspect
 import json
+import sys
 import time
 from functools import partial
 from typing import Generator, Optional
@@ -156,11 +157,36 @@ def test_class_methods():
     # |  `- 0.500 FakeClock.sleep  test/fake_time_util.py:19
     # |- 0.250 ClassWithMethods.long_class_method  test/test_profiler.py:36
     # |  `- 0.250 FakeClock.sleep  test/fake_time_util.py:19
-    # `- 0.250 long_static_method  test/test_profiler.py:32
+    # `- 0.250 ClassWithMethods.long_static_method  test/test_profiler.py:32
     #    `- 0.250 FakeClock.sleep  test/fake_time_util.py:19
     assert text_output.count("0.500 ClassWithMethods.long_method") == 1
     assert text_output.count("0.250 ClassWithMethods.long_class_method") == 1
-    assert text_output.count("0.250 long_static_method") == 1
+    if sys.version_info >= (3, 11):
+        assert text_output.count("0.250 ClassWithMethods.long_static_method") == 1
+    else:
+        assert text_output.count("0.250 long_static_method") == 1
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires code.co_qualname")
+def test_inherited_class_methods():
+    class ChildA(ClassWithMethods):
+        pass
+
+    class ChildB(ClassWithMethods):
+        pass
+
+    profiler = Profiler()
+    with fake_time():
+        with profiler:
+            for obj in [ChildA(), ChildB()]:
+                obj.long_method()
+                obj.long_class_method()
+
+    text_output = profiler.output_text()
+    assert text_output.count("0.500 ClassWithMethods.long_method") == 1
+    assert text_output.count("0.500 ClassWithMethods.long_class_method") == 1
+    assert "ChildA" not in text_output
+    assert "ChildB" not in text_output
 
 
 def test_context_manager():