Add support for displaying numpy scalars (#752)

Author: peytondmurray

.github/workflows/build.yml

@@ -43,6 +43,8 @@ jobs:
       - uses: actions/setup-python@v4
         with:
           python-version: 3.9
+      - name: Install optional numpy dependency
+        run: pip install numpy>=2
       - name: Test
         id: test
         continue-on-error: true

pyproject.toml

@@ -15,6 +15,9 @@ classifiers = [
 ]
 dynamic = ["version"]
 
+[project.optional-dependencies]
+test = ["numpy"]
+
 [tool.maturin]
 bindings = "bin"
 

src/python_data_access.rs

@@ -472,13 +472,53 @@ where
         format!("{}", value.ob_fval)
     } else if value_type_name == "NoneType" {
         "None".to_owned()
+    } else if value_type_name.starts_with("numpy.") {
+        match value_type_name {
+            "numpy.bool" => format_obval::<bool, P>(addr, process)?,
+            "numpy.uint8" => format_obval::<u8, P>(addr, process)?,
+            "numpy.uint16" => format_obval::<u16, P>(addr, process)?,
+            "numpy.uint32" => format_obval::<u32, P>(addr, process)?,
+            "numpy.uint64" => format_obval::<u64, P>(addr, process)?,
+            "numpy.int8" => format_obval::<i8, P>(addr, process)?,
+            "numpy.int16" => format_obval::<i16, P>(addr, process)?,
+            "numpy.int32" => format_obval::<i32, P>(addr, process)?,
+            "numpy.int64" => format_obval::<i64, P>(addr, process)?,
+            "numpy.float32" => format_obval::<f32, P>(addr, process)?,
+            "numpy.float64" => format_obval::<f64, P>(addr, process)?,
+            _ => format!("<{} at 0x{:x}>", value_type_name, addr),
+        }
     } else {
         format!("<{} at 0x{:x}>", value_type_name, addr)
     };
 
     Ok(formatted)
 }
 
+/// Format the numpy scalar to a string.
+///
+/// All numpy scalars have shape:
+/// {
+///     ob_base: PyObject,
+///     obval: <value>,
+/// }
+///
+/// Where `obval` can be of different sizes depending on the scalar type.
+/// We match the size to the value_type_name for this purpose, avoiding the
+/// need to build bindings for the numpy C API.
+///
+/// * `addr`: Address of the numpy scalar
+/// * `process`: Process memory in which the object resides
+fn format_obval<T, P>(addr: usize, process: &P) -> Result<String, Error>
+where
+    T: std::fmt::Display,
+    P: ProcessMemory,
+{
+    let base_addr = addr as *mut u32;
+    let offset = std::mem::size_of::<crate::python_bindings::v3_7_0::PyObject>() as isize;
+    let result = unsafe { process.copy_pointer(base_addr.byte_offset(offset) as *const T)? };
+    Ok(format!("{}", result))
+}
+
 #[cfg(test)]
 pub mod tests {
     // the idea here is to create various cpython interpretator structs locally

tests/integration_test.rs

@@ -280,7 +280,7 @@ fn test_local_vars() {
     let frame = &trace.frames[0];
     let locals = frame.locals.as_ref().unwrap();
 
-    assert_eq!(locals.len(), 9);
+    assert_eq!(locals.len(), 27);
 
     let arg1 = &locals[0];
     assert_eq!(arg1.name, "arg1");
@@ -326,6 +326,95 @@ fn test_local_vars() {
     assert_eq!(local6.name, "local6");
     assert!(!local6.arg);
 
+    // Numpy scalars
+    let local7 = &locals[9];
+    assert_eq!(local7.name, "local7");
+    assert_eq!(local7.repr, Some("true".to_string()));
+
+    let local8 = &locals[10];
+    assert_eq!(local8.name, "local8");
+    assert_eq!(local8.repr, Some("2".to_string()));
+
+    let local9 = &locals[11];
+    assert_eq!(local9.name, "local9");
+    assert_eq!(local9.repr, Some("3".to_string()));
+
+    let local10 = &locals[12];
+    assert_eq!(local10.name, "local10");
+    assert_eq!(local10.repr, Some("42".to_string()));
+
+    let local11 = &locals[13];
+    assert_eq!(local11.name, "local11");
+    assert_eq!(local11.repr, Some("43".to_string()));
+
+    let local12 = &locals[14];
+    assert_eq!(local12.name, "local12");
+    assert_eq!(local12.repr, Some("44".to_string()));
+
+    let local13 = &locals[15];
+    assert_eq!(local13.name, "local13");
+    assert_eq!(local13.repr, Some("45".to_string()));
+
+    let local14 = &locals[16];
+    assert_eq!(local14.name, "local14");
+    assert_eq!(local14.repr, Some("46".to_string()));
+
+    let local15 = &locals[17];
+    assert_eq!(local15.name, "local15");
+    assert_eq!(local15.repr, Some("7".to_string()));
+
+    let local16 = &locals[18];
+    assert_eq!(local16.name, "local16");
+    assert_eq!(local16.repr, Some("8".to_string()));
+
+    fn test_repr_prefix(local: &py_spy::stack_trace::LocalVariable, expected: &str) {
+        assert!(
+            local
+                .repr
+                .as_ref()
+                .map(|result| result.starts_with(expected))
+                .unwrap_or(false),
+            "local '{}' repr = '{:?}' doesn't start with '{}'",
+            &local.name,
+            &local.repr,
+            expected
+        );
+    }
+
+    let local17 = &locals[19];
+    assert_eq!(local17.name, "local17");
+
+    #[cfg(not(windows))]
+    test_repr_prefix(local17, "<numpy.ulonglong at");
+
+    let local18 = &locals[20];
+    assert_eq!(local18.name, "local18");
+    test_repr_prefix(local18, "<numpy.float16 at");
+
+    let local19 = &locals[21];
+    assert_eq!(local19.name, "local19");
+    assert_eq!(local19.repr, Some("0.5".to_string()));
+
+    let local20 = &locals[22];
+    assert_eq!(local20.name, "local20");
+    assert_eq!(local20.repr, Some("0.7".to_string()));
+
+    let local21 = &locals[23];
+    assert_eq!(local21.name, "local21");
+    test_repr_prefix(local21, "<numpy.longdouble at");
+
+    let local22 = &locals[24];
+    assert_eq!(local22.name, "local22");
+    test_repr_prefix(local22, "<numpy.complex64 at");
+
+    let local23 = &locals[25];
+    assert_eq!(local23.name, "local23");
+    test_repr_prefix(local23, "<numpy.complex128 at");
+
+    let local24 = &locals[26];
+    assert_eq!(local24.name, "local24");
+    test_repr_prefix(local24, "<numpy.clongdouble at");
+
     // we only support dictionary lookup on python 3.6+ right now
     if runner.spy.version.major == 3 && runner.spy.version.minor >= 6 {
         assert_eq!(

tests/scripts/local_vars.py

@@ -1,4 +1,5 @@
 import time
+import numpy as np
 
 
 def local_variable_lookup(arg1="foo", arg2=None, arg3=True):
@@ -9,6 +10,36 @@ def local_variable_lookup(arg1="foo", arg2=None, arg3=True):
     local5 = {"a": False, "b": (1, 2, 3)}
     # https://github.com/benfred/py-spy/issues/224
     local6 = ("-" * 115, {"key": {"key": {"key": "value"}}})
+
+    # Numpy scalars
+    # integers
+    local7 = np.bool(True)
+    local8 = np.byte(2)
+
+    local9 = np.int8(3)
+    local10 = np.int16(42)
+    local11 = np.int32(43)
+    local12 = np.int64(44)
+
+    local13 = np.uint8(45)
+    local14 = np.uint16(46)
+    local15 = np.uint32(7)
+    local16 = np.uint64(8)
+
+    local17 = np.ulonglong(11)
+
+    # Floats
+    local18 = np.float16(0.3)
+    local19 = np.float32(0.5)
+    local20 = np.float64(0.7)
+    local21 = np.longdouble(0.9)
+
+    # Complex
+    local22 = np.complex64(0.3+5j)
+    local23 = np.complex128(0.3+5j)
+    local24 = np.clongdouble(0.3+5j)
+
+
     time.sleep(100000)