thread.setname and context (#6726)

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<Self>, vm: &VirtualMachine) -> Vec<PyObjectRef> {
+            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<Self>, vm: &VirtualMachine) -> PyResult {
+            let vars = zelf.borrow_vars();
+            let keys: Vec<PyObjectRef> = vars.keys().map(|k| k.clone().into()).collect();
+            let list = vm.ctx.new_list(keys);
+            <PyList as Iterable>::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<Self>) -> PyRef<Self> {
+            zelf
+        }
+
+        #[pymethod]
+        fn __exit__(
+            zelf: &Py<Self>,
+            _ty: PyObjectRef,
+            _val: PyObjectRef,
+            _tb: PyObjectRef,
+            vm: &VirtualMachine,
+        ) -> PyResult<()> {
+            ContextVar::reset(&zelf.var, zelf.to_owned(), vm)
+        }
     }
 
     impl Constructor for ContextToken {

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)]