diff --git a/crates/stdlib/src/contextvars.rs b/crates/stdlib/src/contextvars.rs index 329342fe6dc..5608b39d1e0 100644 --- a/crates/stdlib/src/contextvars.rs +++ b/crates/stdlib/src/contextvars.rs @@ -24,12 +24,12 @@ thread_local! { mod _contextvars { use crate::vm::{ AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func, - builtins::{PyGenericAlias, PyStrRef, PyType, PyTypeRef}, + builtins::{PyGenericAlias, PyList, PyStrRef, PyType, PyTypeRef}, class::StaticType, common::hash::PyHash, function::{ArgCallable, FuncArgs, OptionalArg}, protocol::{PyMappingMethods, PySequenceMethods}, - types::{AsMapping, AsSequence, Constructor, Hashable, Representable}, + types::{AsMapping, AsSequence, Constructor, Hashable, Iterable, Representable}, }; use core::{ cell::{Cell, RefCell, UnsafeCell}, @@ -163,7 +163,7 @@ mod _contextvars { } } - #[pyclass(with(Constructor, AsMapping, AsSequence))] + #[pyclass(with(Constructor, AsMapping, AsSequence, Iterable))] impl PyContext { #[pymethod] fn run( @@ -205,11 +205,6 @@ mod _contextvars { self.borrow_vars().len() } - #[pymethod] - fn __iter__(&self) -> PyResult { - unimplemented!("Context.__iter__ is currently under construction") - } - #[pymethod] fn get( &self, @@ -238,6 +233,15 @@ mod _contextvars { let vars = zelf.borrow_vars(); vars.values().map(|value| value.to_owned()).collect() } + + // TODO: wrong return type + #[pymethod] + fn items(zelf: PyRef, vm: &VirtualMachine) -> Vec { + let vars = zelf.borrow_vars(); + vars.iter() + .map(|(k, v)| vm.ctx.new_tuple(vec![k.clone().into(), v.clone()]).into()) + .collect() + } } impl Constructor for PyContext { @@ -281,6 +285,15 @@ mod _contextvars { } } + impl Iterable for PyContext { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let vars = zelf.borrow_vars(); + let keys: Vec = vars.keys().map(|k| k.clone().into()).collect(); + let list = vm.ctx.new_list(keys); + ::iter(list, vm) + } + } + #[pyattr] #[pyclass(name, traverse)] #[derive(PyPayload)] @@ -574,6 +587,22 @@ mod _contextvars { ) -> PyGenericAlias { PyGenericAlias::from_args(cls, args, vm) } + + #[pymethod] + fn __enter__(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod] + fn __exit__( + zelf: &Py, + _ty: PyObjectRef, + _val: PyObjectRef, + _tb: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + ContextVar::reset(&zelf.var, zelf.to_owned(), vm) + } } impl Constructor for ContextToken { diff --git a/crates/vm/src/stdlib/thread.rs b/crates/vm/src/stdlib/thread.rs index 0e14fe0e4d1..d51d78015d6 100644 --- a/crates/vm/src/stdlib/thread.rs +++ b/crates/vm/src/stdlib/thread.rs @@ -9,7 +9,7 @@ pub(crate) use _thread::{ pub(crate) mod _thread { use crate::{ AsObject, Py, PyPayload, PyRef, PyResult, VirtualMachine, - builtins::{PyDictRef, PyStr, PyTupleRef, PyType, PyTypeRef}, + builtins::{PyDictRef, PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef}, frame::FrameRef, function::{ArgCallable, Either, FuncArgs, KwArgs, OptionalArg, PySetterValue}, types::{Constructor, GetAttr, Representable, SetAttr}, @@ -260,6 +260,11 @@ pub(crate) mod _thread { Ok(()) } + #[pymethod] + fn locked(&self) -> bool { + self.mu.is_locked() + } + #[pymethod] fn _is_owned(&self) -> bool { self.mu.is_owned_by_current_thread() @@ -293,6 +298,47 @@ pub(crate) mod _thread { current_thread_id() } + /// Set the name of the current thread + #[pyfunction] + fn set_name(name: PyStrRef) { + #[cfg(target_os = "linux")] + { + use std::ffi::CString; + if let Ok(c_name) = CString::new(name.as_str()) { + // pthread_setname_np on Linux has a 16-byte limit including null terminator + // TODO: Potential UTF-8 boundary issue when truncating thread name on Linux. + // https://github.com/RustPython/RustPython/pull/6726/changes#r2689379171 + let truncated = if c_name.as_bytes().len() > 15 { + CString::new(&c_name.as_bytes()[..15]).unwrap_or(c_name) + } else { + c_name + }; + unsafe { + libc::pthread_setname_np(libc::pthread_self(), truncated.as_ptr()); + } + } + } + #[cfg(target_os = "macos")] + { + use std::ffi::CString; + if let Ok(c_name) = CString::new(name.as_str()) { + unsafe { + libc::pthread_setname_np(c_name.as_ptr()); + } + } + } + #[cfg(windows)] + { + // Windows doesn't have a simple pthread_setname_np equivalent + // SetThreadDescription requires Windows 10+ + let _ = name; + } + #[cfg(not(any(target_os = "linux", target_os = "macos", windows)))] + { + let _ = name; + } + } + /// Get OS-level thread ID (pthread_self on Unix) /// This is important for fork compatibility - the ID must remain stable after fork #[cfg(unix)]