Merge pull request #332 from joerick/del_cls

Fix #331 - prevent a crash when a cls or self variable is del'd

pyinstrument/low_level/stat_profile.c

@@ -310,19 +310,31 @@ _get_class_name_of_frame(PyFrameObject *frame, PyCodeObject *code) {
         return NULL;
     }
 
-    if (has_self) {
+    // 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) {
-            result = _PyType_Name(self->ob_type);
+
+        if (!self) {
+            PyErr_Clear();
+            Py_DECREF(locals);
+            return NULL;
         }
+
+        result = _PyType_Name(self->ob_type);
     }
-    else if (has_cls) {
+    else if (has_cls && PyMapping_HasKey(locals, CLS_STRING)) {
         PyObject *cls = PyObject_GetItem(locals, CLS_STRING);
-        if (cls) {
-            if (PyType_Check(cls)) {
-                PyTypeObject *type = (PyTypeObject *)cls;
-                result = _PyType_Name(type);
-            }
+
+        if (!cls) {
+            PyErr_Clear();
+            Py_DECREF(locals);
+            return NULL;
+        }
+
+        if (PyType_Check(cls)) {
+            PyTypeObject *type = (PyTypeObject *)cls;
+            result = _PyType_Name(type);
         }
     }
 

test/low_level/test_frame_info.py

@@ -1,28 +1,45 @@
 import inspect
 
+import pytest
+
 from pyinstrument.low_level import stat_profile as stat_profile_c
 from pyinstrument.low_level import stat_profile_python
 
 
 class AClass:
-    def get_frame_info_for_a_method(self, getter_function):
+    def get_frame_info_for_a_method(self, getter_function, del_local):
+        if del_local:
+            del self
         frame = inspect.currentframe()
         assert frame
         return getter_function(frame)
 
-    def get_frame_info_with_cell_variable(self, getter_function):
-        frame = inspect.currentframe()
-        assert frame
-
+    def get_frame_info_with_cell_variable(self, getter_function, del_local):
         def an_inner_function():
             # reference self to make it a cell variable
             if self:
                 pass
 
+        if del_local:
+            del self
+        frame = inspect.currentframe()
+        assert frame
+
         return getter_function(frame)
 
     @classmethod
-    def get_frame_info_for_a_class_method(cls, getter_function):
+    def get_frame_info_for_a_class_method(cls, getter_function, del_local):
+        if del_local:
+            del cls
+        frame = inspect.currentframe()
+        assert frame
+        return getter_function(frame)
+
+    @classmethod
+    def get_frame_info_for_a_class_method_where_cls_is_reassigned(cls, getter_function, del_local):
+        cls = 1
+        if del_local:
+            del cls
         frame = inspect.currentframe()
         assert frame
         return getter_function(frame)
@@ -59,17 +76,21 @@ def test_frame_info_hide_false():
     assert stat_profile_c.get_frame_info(frame) == stat_profile_python.get_frame_info(frame)
 
 
-def test_frame_info_with_classes():
-    instance = AClass()
+instance = AClass()
 
-    test_functions = [
+
+@pytest.mark.parametrize(
+    "test_function",
+    [
         instance.get_frame_info_for_a_method,
         AClass.get_frame_info_for_a_class_method,
         instance.get_frame_info_with_cell_variable,
-    ]
-
-    for test_function in test_functions:
-        c_frame_info = test_function(stat_profile_c.get_frame_info)
-        py_frame_info = test_function(stat_profile_python.get_frame_info)
-
-        assert c_frame_info == py_frame_info
+        AClass.get_frame_info_for_a_class_method_where_cls_is_reassigned,
+    ],
+)
+@pytest.mark.parametrize("del_local", [True, False])
+def test_frame_info_with_classes(test_function, del_local):
+    c_frame_info = test_function(stat_profile_c.get_frame_info, del_local=del_local)
+    py_frame_info = test_function(stat_profile_python.get_frame_info, del_local=del_local)
+
+    assert c_frame_info == py_frame_info