@@ -12,17 +12,22 @@ thread_local! {
1212#[ pymodule]
1313mod _contextvars {
1414 use crate :: vm:: {
15- AsObject , Py , PyObjectRef , PyPayload , PyRef , PyResult , VirtualMachine , atomic_func,
16- builtins:: { PyGenericAlias , PyList , PyStr , PyType , PyTypeRef } ,
15+ AsObject , Py , PyObject , PyObjectRef , PyPayload , PyRef , PyResult , VirtualMachine ,
16+ atomic_func,
17+ builtins:: { PyGenericAlias , PyList , PyStr , PyStrRef , PyType , PyTypeRef } ,
1718 class:: StaticType ,
19+ class_or_notimplemented,
1820 common:: {
1921 hash:: PyHash ,
2022 lock:: { LazyLock , PyMutex } ,
2123 wtf8:: Wtf8Buf ,
2224 } ,
23- function:: { FuncArgs , OptionalArg } ,
25+ function:: { FuncArgs , OptionalArg , PyComparisonValue } ,
2426 protocol:: { PyMappingMethods , PySequenceMethods } ,
25- types:: { AsMapping , AsSequence , Constructor , Hashable , Iterable , Representable } ,
27+ types:: {
28+ AsMapping , AsSequence , Comparable , Constructor , Hashable , Iterable , PyComparisonOp ,
29+ Representable ,
30+ } ,
2631 } ;
2732 use core:: sync:: atomic:: { AtomicBool , AtomicI64 , AtomicUsize , Ordering } ;
2833 use indexmap:: IndexMap ;
@@ -56,7 +61,7 @@ mod _contextvars {
5661 }
5762
5863 #[ pyattr]
59- #[ pyclass( name = "Context" ) ]
64+ #[ pyclass( name = "Context" , unhashable = true ) ]
6065 #[ derive( Debug , PyPayload ) ]
6166 pub ( crate ) struct PyContext {
6267 // not to confuse with vm::Context
@@ -166,7 +171,7 @@ mod _contextvars {
166171 }
167172 }
168173
169- #[ pyclass( with( Constructor , AsMapping , AsSequence , Iterable ) ) ]
174+ #[ pyclass( with( Constructor , AsMapping , AsSequence , Iterable , Comparable ) ) ]
170175 impl PyContext {
171176 #[ pymethod]
172177 fn run ( zelf : & Py < Self > , args : FuncArgs , vm : & VirtualMachine ) -> PyResult {
@@ -307,12 +312,51 @@ mod _contextvars {
307312 }
308313 }
309314
315+ impl Comparable for PyContext {
316+ fn cmp (
317+ zelf : & Py < Self > ,
318+ other : & PyObject ,
319+ op : PyComparisonOp ,
320+ vm : & VirtualMachine ,
321+ ) -> PyResult < PyComparisonValue > {
322+ let other = class_or_notimplemented ! ( Self , other) ;
323+ op. eq_only ( || Ok ( vars_eq ( zelf, other, vm) ?. into ( ) ) )
324+ }
325+ }
326+
327+ fn vars_eq ( left : & Py < PyContext > , right : & Py < PyContext > , vm : & VirtualMachine ) -> PyResult < bool > {
328+ if left. is ( right) {
329+ return Ok ( true ) ;
330+ }
331+ // Snapshot first: value __eq__ may re-enter Context.run / ContextVar.set.
332+ let pairs: Vec < ( PyObjectRef , PyObjectRef ) > = {
333+ let left_vars = left. borrow_vars ( ) ;
334+ let right_vars = right. borrow_vars ( ) ;
335+ if left_vars. len ( ) != right_vars. len ( ) {
336+ return Ok ( false ) ;
337+ }
338+ let mut pairs = Vec :: with_capacity ( left_vars. len ( ) ) ;
339+ for ( key, left_value) in left_vars. iter ( ) {
340+ let Some ( right_value) = right_vars. get ( key) else {
341+ return Ok ( false ) ;
342+ } ;
343+ pairs. push ( ( left_value. clone ( ) , right_value. clone ( ) ) ) ;
344+ }
345+ pairs
346+ } ;
347+ for ( left_value, right_value) in pairs {
348+ if !vm. bool_eq ( & left_value, & right_value) ? {
349+ return Ok ( false ) ;
350+ }
351+ }
352+ Ok ( true )
353+ }
354+
310355 #[ pyattr]
311356 #[ pyclass( name, traverse) ]
312357 #[ derive( PyPayload ) ]
313358 struct ContextVar {
314- #[ pytraverse( skip) ]
315- name : String ,
359+ name : PyStrRef ,
316360 default : Option < PyObjectRef > ,
317361 #[ pytraverse( skip) ]
318362 cached : PyMutex < Option < ContextVarCache > > ,
@@ -390,7 +434,7 @@ mod _contextvars {
390434 #[ pyclass( with( Constructor , Hashable , Representable ) ) ]
391435 impl ContextVar {
392436 #[ pygetset]
393- fn name ( & self ) -> String {
437+ fn name ( & self ) -> PyStrRef {
394438 self . name . clone ( )
395439 }
396440
@@ -524,7 +568,6 @@ mod _contextvars {
524568 . downcast :: < PyStr > ( )
525569 . map_err ( |_| vm. new_type_error ( "context variable name must be a str" ) ) ?;
526570 let name_hash = name. as_object ( ) . hash ( vm) ?;
527- let name = name. to_string ( ) ;
528571
529572 let var = Self {
530573 name,
@@ -559,17 +602,18 @@ mod _contextvars {
559602 }
560603
561604 impl Representable for ContextVar {
562- #[ inline]
563- fn repr_str ( zelf : & Py < Self > , vm : & VirtualMachine ) -> PyResult < String > {
564- let name = zelf. name . as_str ( ) ;
605+ fn repr_wtf8 ( zelf : & Py < Self > , vm : & VirtualMachine ) -> PyResult < Wtf8Buf > {
606+ let name = zelf. name . as_object ( ) . repr ( vm) ?;
565607 let id = zelf. get_id ( ) ;
566-
567- Ok ( if let Some ( arg) = zelf. default . as_ref ( ) {
568- let default = arg. str ( vm) . ok ( ) ;
569- format ! ( "<ContextVar name='{name}' default={default:?} at {id:#x}>" , )
570- } else {
571- format ! ( "<ContextVar name='{name}' at {id:#x}>" )
572- } )
608+ let mut result = Wtf8Buf :: from ( "<ContextVar name=" ) ;
609+ result. push_wtf8 ( name. as_wtf8 ( ) ) ;
610+ if let Some ( arg) = zelf. default . as_ref ( ) {
611+ let default = arg. repr ( vm) ?;
612+ result. push_str ( " default=" ) ;
613+ result. push_wtf8 ( default. as_wtf8 ( ) ) ;
614+ }
615+ result. push_str ( & format ! ( " at {id:#x}>" ) ) ;
616+ Ok ( result)
573617 }
574618 }
575619
0 commit comments