Add --class-name param to prepend class names to methods

Regular methods and classmethods supported. No support for static methods.
This is because the frame object doesn't refer to the function, only its
code, which leaves no other option than to try to find `self` or `cls`
in locals. The alternative of using the `gc` module is horrendously slow
even in the same process and also impractical in a remote process.

Author: karolinepauls

src/config.rs

@@ -56,6 +56,8 @@ pub struct Config {
     #[doc(hidden)]
     pub lineno: LineNo,
     #[doc(hidden)]
+    pub include_class_name: bool,
+    #[doc(hidden)]
     pub refresh_seconds: f64,
     #[doc(hidden)]
     pub core_filename: Option<String>,
@@ -137,6 +139,7 @@ impl Default for Config {
             subprocesses: false,
             full_filenames: false,
             lineno: LineNo::LastInstruction,
+            include_class_name: false,
             refresh_seconds: 1.0,
             core_filename: None,
         }
@@ -192,6 +195,9 @@ impl Config {
         let full_filenames = Arg::new("full_filenames").long("full-filenames").help(
             "Show full Python filenames, instead of shortening to show only the package part",
         );
+        let include_class_name = Arg::new("include_class_name").long("class-name").help(
+            "Prepend class name to method names as long as the `self`/`cls` argument naming convention is followed (doesn't work for staticmethods)",
+        );
         let program = Arg::new("python_program")
             .help("commandline of a python program to run")
             .multiple_values(true);
@@ -219,6 +225,7 @@ impl Config {
             .arg(program.clone())
             .arg(pid.clone().required_unless_present("python_program"))
             .arg(full_filenames.clone())
+            .arg(include_class_name.clone())
             .arg(
                 Arg::new("output")
                     .short('o')
@@ -286,6 +293,7 @@ impl Config {
             .arg(rate.clone())
             .arg(subprocesses.clone())
             .arg(full_filenames.clone())
+            .arg(include_class_name.clone())
             .arg(gil.clone())
             .arg(idle.clone())
             .arg(top_delay.clone());
@@ -311,6 +319,7 @@ impl Config {
         );
 
         let dump = dump.arg(full_filenames.clone())
+            .arg(include_class_name.clone())
             .arg(Arg::new("locals")
                 .short('l')
                 .long("locals")
@@ -442,6 +451,8 @@ impl Config {
         });
 
         config.full_filenames = matches.occurrences_of("full_filenames") > 0;
+        config.include_class_name = matches.occurrences_of("include_class_name") > 0;
+
         if cfg!(feature = "unwind") {
             config.native = matches.occurrences_of("native") > 0;
         }

src/python_data_access.rs

@@ -353,6 +353,25 @@ const PY_TPFLAGS_BYTES_SUBCLASS: usize = 1 << 27;
 const PY_TPFLAGS_STRING_SUBCLASS: usize = 1 << 28;
 const PY_TPFLAGS_DICT_SUBCLASS: usize = 1 << 29;
 
+const MAX_TYPE_NAME_LEN: usize = 64;
+
+/// Get the type's name (truncating to MAX_TYPE_NAME_LEN bytes if longer)
+pub fn extract_type_name<T, P>(process: &P, value_type: &T) -> Result<String, Error>
+where
+    T: TypeObject,
+    P: ProcessMemory,
+{
+    let mut value_type_name = process.copy(value_type.name() as usize, MAX_TYPE_NAME_LEN)?;
+    let length = value_type_name
+        .iter()
+        .position(|&x| x == 0)
+        .unwrap_or(MAX_TYPE_NAME_LEN);
+    value_type_name.truncate(length);
+
+    let string = String::from_utf8(value_type_name)?;
+    Ok(string)
+}
+
 /// Converts a python variable in the other process to a human readable string
 pub fn format_variable<I, P>(
     process: &P,
@@ -373,14 +392,7 @@ where
     let value: I::Object = process.copy_struct(addr)?;
     let value_type = process.copy_pointer(value.ob_type())?;
 
-    // get the typename (truncating to 128 bytes if longer)
-    let max_type_len = 128;
-    let value_type_name = process.copy(value_type.name() as usize, max_type_len)?;
-    let length = value_type_name
-        .iter()
-        .position(|&x| x == 0)
-        .unwrap_or(max_type_len);
-    let value_type_name = std::str::from_utf8(&value_type_name[..length])?;
+    let value_type_name = extract_type_name(process, &value_type)?;
 
     let format_int = |value: i64| {
         if value_type_name == "bool" {
@@ -476,7 +488,7 @@ where
     } else if value_type_name == "NoneType" {
         "None".to_owned()
     } else if value_type_name.starts_with("numpy.") {
-        match value_type_name {
+        match value_type_name.as_str() {
             "numpy.bool" => format_obval::<bool, P>(addr, process)?,
             "numpy.uint8" => format_obval::<u8, P>(addr, process)?,
             "numpy.uint16" => format_obval::<u16, P>(addr, process)?,

src/python_spy.rs

@@ -250,11 +250,12 @@ impl PythonSpy {
                 continue;
             }
 
-            let mut trace = get_stack_trace(
+            let mut trace = get_stack_trace::<I, <I as InterpreterState>::ThreadState, Process>(
                 &thread,
                 &self.process,
                 self.config.dump_locals > 0,
                 self.config.lineno,
+                self.config.include_class_name,
             )?;
 
             // Try getting the native thread id

src/stack_trace.rs

@@ -6,9 +6,9 @@ use remoteprocess::{Pid, ProcessMemory};
 use serde_derive::Serialize;
 
 use crate::config::{Config, LineNo};
-use crate::python_data_access::{copy_bytes, copy_string};
+use crate::python_data_access::{copy_bytes, copy_string, extract_type_name};
 use crate::python_interpreters::{
-    CodeObject, FrameObject, InterpreterState, ThreadState, TupleObject,
+    CodeObject, FrameObject, InterpreterState, Object, ThreadState, TupleObject, TypeObject,
 };
 
 /// Call stack for a single python thread
@@ -68,6 +68,8 @@ pub struct ProcessInfo {
     pub parent: Option<Box<ProcessInfo>>,
 }
 
+const PY_TPFLAGS_TYPE_SUBCLASS: usize = 1 << 31;
+
 /// Given an InterpreterState, this function returns a vector of stack traces for each thread
 pub fn get_stack_traces<I, P>(
     interpreter_address: usize,
@@ -90,13 +92,20 @@ where
 
     let lineno = config.map(|c| c.lineno).unwrap_or(LineNo::NoLine);
     let dump_locals = config.map(|c| c.dump_locals).unwrap_or(0);
+    let include_class_name = config.map(|c| c.include_class_name).unwrap_or(false);
 
     while !threads.is_null() {
         let thread = process
             .copy_pointer(threads)
             .context("Failed to copy PyThreadState")?;
 
-        let mut trace = get_stack_trace(&thread, process, dump_locals > 0, lineno)?;
+        let mut trace = get_stack_trace::<I, <I as InterpreterState>::ThreadState, P>(
+            &thread,
+            process,
+            dump_locals > 0,
+            lineno,
+            include_class_name,
+        )?;
         trace.owns_gil = trace.thread_id == gil_thread_id;
 
         ret.push(trace);
@@ -110,13 +119,15 @@ where
 }
 
 /// Gets a stack trace for an individual thread
-pub fn get_stack_trace<T, P>(
+pub fn get_stack_trace<I, T, P>(
     thread: &T,
     process: &P,
     copy_locals: bool,
     lineno: LineNo,
+    include_class_name: bool,
 ) -> Result<StackTrace, Error>
 where
+    I: InterpreterState,
     T: ThreadState,
     P: ProcessMemory,
 {
@@ -165,7 +176,7 @@ where
             continue;
         }
         let filename = filename?;
-        let name = name?;
+        let mut name = name?;
 
         // skip <shim> entries in python 3.12+
         // Unset file/function name in py3.13 means this is a shim.
@@ -195,11 +206,27 @@ where
         };
 
         let locals = if copy_locals {
-            Some(get_locals(&code, frame_ptr, &frame, process)?)
+            Some(get_locals(&code, frame_ptr, &frame, process, false)?)
         } else {
             None
         };
 
+        // cls or self are always the first argument in methods and classmethods (and first local)
+        if include_class_name && code.argcount() > 0 {
+            let found_locals = match locals {
+                Some(ref found_locals) => found_locals,
+                None => &get_locals(&code, frame_ptr, &frame, process, true)?,
+            };
+            // Errors copying self/cls type name are silenced.
+            if let Some(class_name) = found_locals
+                .get(0)
+                .and_then(|first_arg| get_class_name_from_arg::<I, P>(process, first_arg).ok())
+                .flatten()
+            {
+                name = format!("{}.{}", class_name, name);
+            }
+        }
+
         let is_entry = frame.is_entry();
 
         frames.push(Frame {
@@ -273,6 +300,7 @@ fn get_locals<C: CodeObject, F: FrameObject, P: ProcessMemory>(
     frameptr: *const F,
     frame: &F,
     process: &P,
+    first_var_only: bool,
 ) -> Result<Vec<LocalVariable>, Error> {
     let local_count = code.nlocals() as usize;
     let argcount = code.argcount() as usize;
@@ -283,7 +311,12 @@ fn get_locals<C: CodeObject, F: FrameObject, P: ProcessMemory>(
 
     let mut ret = Vec::new();
 
-    for i in 0..local_count {
+    let vars_to_copy = if first_var_only {
+        std::cmp::min(local_count, 1)
+    } else {
+        local_count
+    };
+    for i in 0..vars_to_copy {
         let nameptr: *const C::StringObject =
             process.copy_struct(varnames.address(code.varnames() as usize, i))?;
         let name = copy_string(nameptr, process)?;
@@ -301,6 +334,49 @@ fn get_locals<C: CodeObject, F: FrameObject, P: ProcessMemory>(
     Ok(ret)
 }
 
+/// Get class from a `self` or `cls` argument, as long as its type matches expectations.
+fn get_class_name_from_arg<I, P>(
+    process: &P,
+    first_local: &LocalVariable,
+) -> Result<Option<String>, Error>
+where
+    I: InterpreterState,
+    P: ProcessMemory,
+{
+    // If the first variable isn't an argument, there are no arguments, so the fn isn't a normal
+    // method or a class method.
+    if !first_local.arg {
+        return Ok(None);
+    }
+
+    let first_arg_name = &first_local.name;
+    if first_arg_name != "self" && first_arg_name != "cls" {
+        return Ok(None);
+    }
+
+    let value: I::Object = process.copy_struct(first_local.addr)?;
+    let mut value_type = process.copy_pointer(value.ob_type())?;
+    let is_type = value_type.flags() & PY_TPFLAGS_TYPE_SUBCLASS != 0;
+
+    // validate that the first argument is:
+    // - an instance of something else than `type` if it is called "self"
+    // - an instance of `type` if it is called "cls"
+    match (first_arg_name.as_str(), is_type) {
+        ("self", false) => {}
+        ("cls", true) => {
+            // Copy the remote argument struct, but this time as PyTypeObject, rather than
+            // PyObject. We needed to read type flags from PyObject to know that this is a
+            // PyTypeObject.
+            value_type = process.copy_struct(first_local.addr)?;
+        }
+        _ => {
+            return Ok(None);
+        }
+    }
+
+    Ok(Some(extract_type_name(process, &value_type)?))
+}
+
 pub fn get_gil_threadid<I: InterpreterState, P: ProcessMemory>(
     threadstate_address: usize,
     process: &P,