From 5afe140abd531eb63aa89bbc50e45efbc0055394 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" <69878+youknowone@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:46:20 +0900 Subject: [PATCH] Store ContextVar.name as the original str object (#8682) * Store ContextVar.name as the original str object Keep the constructor argument instead of converting it to a Rust String so lone surrogates and str subclasses are preserved. Build __repr__ from the name and default reprs. Assisted-by: Grok:4.6 * Compare Context objects by their variable maps Implement Context ==/!= from the mapping contents and mark Context unhashable. Snapshot the map before value equality so re-entrant ContextVar.set during __eq__ cannot deadlock. Enable test_context.ContextTest.test_context_run_4. Assisted-by: Grok:4.6 --- Lib/test/test_context.py | 1 - crates/stdlib/src/contextvars.rs | 84 ++++++++++++++++------ extra_tests/snippets/stdlib_contextvars.py | 56 +++++++++++++++ 3 files changed, 120 insertions(+), 21 deletions(-) create mode 100644 extra_tests/snippets/stdlib_contextvars.py diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py index 145394003..ef20495dc 100644 --- a/Lib/test/test_context.py +++ b/Lib/test/test_context.py @@ -163,7 +163,6 @@ class ContextTest(unittest.TestCase): with self.assertRaises(ZeroDivisionError): ctx.run(func, 1, 2, a=123) - @unittest.expectedFailure # TODO: RUSTPYTHON @isolated_context def test_context_run_4(self): ctx1 = contextvars.Context() diff --git a/crates/stdlib/src/contextvars.rs b/crates/stdlib/src/contextvars.rs index e880625eb..9f28d4d36 100644 --- a/crates/stdlib/src/contextvars.rs +++ b/crates/stdlib/src/contextvars.rs @@ -12,17 +12,22 @@ thread_local! { #[pymodule] mod _contextvars { use crate::vm::{ - AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func, - builtins::{PyGenericAlias, PyList, PyStr, PyType, PyTypeRef}, + AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + atomic_func, + builtins::{PyGenericAlias, PyList, PyStr, PyStrRef, PyType, PyTypeRef}, class::StaticType, + class_or_notimplemented, common::{ hash::PyHash, lock::{LazyLock, PyMutex}, wtf8::Wtf8Buf, }, - function::{FuncArgs, OptionalArg}, + function::{FuncArgs, OptionalArg, PyComparisonValue}, protocol::{PyMappingMethods, PySequenceMethods}, - types::{AsMapping, AsSequence, Constructor, Hashable, Iterable, Representable}, + types::{ + AsMapping, AsSequence, Comparable, Constructor, Hashable, Iterable, PyComparisonOp, + Representable, + }, }; use core::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering}; use indexmap::IndexMap; @@ -56,7 +61,7 @@ mod _contextvars { } #[pyattr] - #[pyclass(name = "Context")] + #[pyclass(name = "Context", unhashable = true)] #[derive(Debug, PyPayload)] pub(crate) struct PyContext { // not to confuse with vm::Context @@ -166,7 +171,7 @@ mod _contextvars { } } - #[pyclass(with(Constructor, AsMapping, AsSequence, Iterable))] + #[pyclass(with(Constructor, AsMapping, AsSequence, Iterable, Comparable))] impl PyContext { #[pymethod] fn run(zelf: &Py, args: FuncArgs, vm: &VirtualMachine) -> PyResult { @@ -307,12 +312,51 @@ mod _contextvars { } } + impl Comparable for PyContext { + fn cmp( + zelf: &Py, + other: &PyObject, + op: PyComparisonOp, + vm: &VirtualMachine, + ) -> PyResult { + let other = class_or_notimplemented!(Self, other); + op.eq_only(|| Ok(vars_eq(zelf, other, vm)?.into())) + } + } + + fn vars_eq(left: &Py, right: &Py, vm: &VirtualMachine) -> PyResult { + if left.is(right) { + return Ok(true); + } + // Snapshot first: value __eq__ may re-enter Context.run / ContextVar.set. + let pairs: Vec<(PyObjectRef, PyObjectRef)> = { + let left_vars = left.borrow_vars(); + let right_vars = right.borrow_vars(); + if left_vars.len() != right_vars.len() { + return Ok(false); + } + let mut pairs = Vec::with_capacity(left_vars.len()); + for (key, left_value) in left_vars.iter() { + let Some(right_value) = right_vars.get(key) else { + return Ok(false); + }; + pairs.push((left_value.clone(), right_value.clone())); + } + pairs + }; + for (left_value, right_value) in pairs { + if !vm.bool_eq(&left_value, &right_value)? { + return Ok(false); + } + } + Ok(true) + } + #[pyattr] #[pyclass(name, traverse)] #[derive(PyPayload)] struct ContextVar { - #[pytraverse(skip)] - name: String, + name: PyStrRef, default: Option, #[pytraverse(skip)] cached: PyMutex>, @@ -390,7 +434,7 @@ mod _contextvars { #[pyclass(with(Constructor, Hashable, Representable))] impl ContextVar { #[pygetset] - fn name(&self) -> String { + fn name(&self) -> PyStrRef { self.name.clone() } @@ -524,7 +568,6 @@ mod _contextvars { .downcast::() .map_err(|_| vm.new_type_error("context variable name must be a str"))?; let name_hash = name.as_object().hash(vm)?; - let name = name.to_string(); let var = Self { name, @@ -559,17 +602,18 @@ mod _contextvars { } impl Representable for ContextVar { - #[inline] - fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let name = zelf.name.as_str(); + fn repr_wtf8(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let name = zelf.name.as_object().repr(vm)?; let id = zelf.get_id(); - - Ok(if let Some(arg) = zelf.default.as_ref() { - let default = arg.str(vm).ok(); - format!("",) - } else { - format!("") - }) + let mut result = Wtf8Buf::from("")); + Ok(result) } } diff --git a/extra_tests/snippets/stdlib_contextvars.py b/extra_tests/snippets/stdlib_contextvars.py new file mode 100644 index 000000000..e0dc0d2c5 --- /dev/null +++ b/extra_tests/snippets/stdlib_contextvars.py @@ -0,0 +1,56 @@ +from contextvars import Context, ContextVar, copy_context + +from testutils import assert_raises + +name = "\ud800" +var = ContextVar(name) +assert var.name == name +assert var.name is name + + +class StrSub(str): + pass + + +sub = StrSub("foo") +var_sub = ContextVar(sub) +assert var_sub.name is sub + +assert_raises(TypeError, ContextVar, 1) + +ctx1 = Context() +ctx2 = Context() +nested = ContextVar("var") + + +def func2(): + assert nested.get(None) is None + + +def func1(): + assert nested.get(None) is None + nested.set("spam") + ctx2.run(func2) + assert nested.get(None) == "spam" + cur = copy_context() + assert len(cur) == 1 + assert cur[nested] == "spam" + return cur + + +returned_ctx = ctx1.run(func1) +assert ctx1 == returned_ctx +assert returned_ctx[nested] == "spam" +assert nested in returned_ctx +assert_raises(TypeError, hash, ctx1) + + +class ReentrantEq: + def __eq__(self, other): + ctx1.run(lambda: nested.set(object())) + return True + + +ctx1.run(nested.set, ReentrantEq()) +ctx2.run(nested.set, object()) +assert ctx1 == ctx2