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

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,
         }
@@ -188,6 +191,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);
@@ -215,6 +221,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')
@@ -282,6 +289,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());
@@ -307,6 +315,7 @@ impl Config {
         );
 
         let dump = dump.arg(full_filenames.clone())
+            .arg(include_class_name.clone())
             .arg(Arg::new("locals")
                 .short('l')
                 .long("locals")
@@ -429,6 +438,7 @@ impl Config {
             .value_of("pid")
             .map(|p| p.parse().expect("invalid pid"));
         config.full_filenames = matches.occurrences_of("full_filenames") > 0;
+        config.include_class_name = matches.occurrences_of("include_class_name") > 0;
         if cfg!(unwind) {
             config.native = matches.occurrences_of("native") > 0;
         }

src/python_data_access.rs

@@ -276,6 +276,24 @@ const PY_TPFLAGS_BYTES_SUBCLASS: usize = 1 << 27;
 const PY_TPFLAGS_STRING_SUBCLASS: usize = 1 << 28;
 const PY_TPFLAGS_DICT_SUBCLASS: usize = 1 << 29;
 
+/// Get the type name (truncating to 128 bytes if longer)
+pub fn extract_type_name<T, P>(process: &P, value_type: &T) -> Result<String, Error>
+where
+    T: TypeObject,
+    P: ProcessMemory,
+{
+    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 string = String::from_utf8(value_type_name[..length].to_vec())?;
+
+    Ok(string)
+}
+
 /// Converts a python variable in the other process to a human readable string
 pub fn format_variable<I, P>(
     process: &P,
@@ -296,14 +314,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" {

src/python_spy.rs

@@ -238,11 +238,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
@@ -66,6 +66,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: &I,
@@ -84,13 +86,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);
@@ -104,13 +113,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,
 {
@@ -133,7 +144,7 @@ where
             .context("Failed to copy PyCodeObject")?;
 
         let filename = copy_string(code.filename(), process).context("Failed to copy filename")?;
-        let name = copy_string(code.name(), process).context("Failed to copy function name")?;
+        let mut name = copy_string(code.name(), process).context("Failed to copy function name")?;
 
         let line = match lineno {
             LineNo::NoLine => 0,
@@ -154,8 +165,28 @@ where
             },
         };
 
-        let locals = if copy_locals {
-            Some(get_locals(&code, frame_ptr, &frame, process)?)
+        // Grab locals, which may be needed for locals display or to find the class if the fn is
+        // a method.
+        let locals = if copy_locals || (include_class_name && code.argcount() > 0) {
+            // Only copy the first local if we only want to find the class name.
+            let first_var_only = !copy_locals;
+            let found_locals = get_locals(&code, frame_ptr, &frame, process, first_var_only)?;
+
+            if include_class_name && !found_locals.is_empty() {
+                let first_arg = &found_locals[0];
+                if let Some(class_name) =
+                    get_class_name_from_arg::<I, P>(process, first_arg)?.as_ref()
+                {
+                    // e.g. 'method' is turned into 'ClassName.method'
+                    name = format!("{}.{}", class_name, name);
+                }
+            }
+
+            if copy_locals {
+                Some(found_locals)
+            } else {
+                None
+            }
         } else {
             None
         };
@@ -229,6 +260,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;
@@ -239,7 +271,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)?;
@@ -257,6 +294,47 @@ 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
+            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,