sys.get_coroutine_origin_tracking_depth

youknowone · youknowone · commit c97d3d3372e0 · 2024-05-01T17:54:42.000+09:00
diff --git a/vm/src/stdlib/sys.rs b/vm/src/stdlib/sys.rs
@@ -692,6 +692,20 @@ mod sys {
         vm.use_tracing.set(tracing);
     }
 
+    #[pyfunction]
+    fn set_coroutine_origin_tracking_depth(depth: i32, vm: &VirtualMachine) -> PyResult<()> {
+        if depth < 0 {
+            return Err(vm.new_value_error("depth must be >= 0".to_owned()));
+        }
+        crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.set(depth as _));
+        Ok(())
+    }
+
+    #[pyfunction]
+    fn get_coroutine_origin_tracking_depth() -> i32 {
+        crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.get()) as _
+    }
+
     /// sys.flags
     ///
     /// Flags provided through command line arguments or environment vars.
diff --git a/vm/src/vm/thread.rs b/vm/src/vm/thread.rs
@@ -1,14 +1,16 @@
 use crate::{AsObject, PyObject, VirtualMachine};
 use itertools::Itertools;
 use std::{
-    cell::RefCell,
-    ptr::{null, NonNull},
+    cell::{Cell, RefCell},
+    ptr::NonNull,
     thread_local,
 };
 
 thread_local! {
     pub(super) static VM_STACK: RefCell<Vec<NonNull<VirtualMachine>>> = Vec::with_capacity(1).into();
-    static VM_CURRENT: RefCell<*const VirtualMachine> = null::<VirtualMachine>().into();
+    static VM_CURRENT: RefCell<*const VirtualMachine> = std::ptr::null::<VirtualMachine>().into();
+
+    pub(crate) static COROUTINE_ORIGIN_TRACKING_DEPTH: Cell<u32> = Cell::new(0);
 }
 
 pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
@@ -142,7 +144,6 @@ impl VirtualMachine {
     /// specific guaranteed behavior.
     #[cfg(feature = "threading")]
     pub fn new_thread(&self) -> ThreadedVirtualMachine {
-        use std::cell::Cell;
         let vm = VirtualMachine {
             builtins: self.builtins.clone(),
             sys_module: self.sys_module.clone(),
