Code nits in contextvars.rs

ShaharNaveh · ShaharNaveh · commit 4a1f65583f3e · 2026-05-11T12:39:51.000+03:00
diff --git a/crates/stdlib/src/contextvars.rs b/crates/stdlib/src/contextvars.rs
@@ -15,7 +15,7 @@ mod _contextvars {
         AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
         builtins::{PyGenericAlias, PyList, PyStrRef, PyType, PyTypeRef},
         class::StaticType,
-        common::{hash::PyHash, wtf8::Wtf8Buf},
+        common::{hash::PyHash, lock::LazyLock, wtf8::Wtf8Buf},
         function::{ArgCallable, FuncArgs, OptionalArg},
         protocol::{PyMappingMethods, PySequenceMethods},
         types::{AsMapping, AsSequence, Constructor, Hashable, Iterable, Representable},
@@ -26,7 +26,6 @@ mod _contextvars {
     };
     use crossbeam_utils::atomic::AtomicCell;
     use indexmap::IndexMap;
-    use rustpython_common::lock::LazyLock;
 
     // TODO: Real hamt implementation
     type Hamt = IndexMap<PyRef<ContextVar>, PyObjectRef, ahash::RandomState>;
@@ -89,11 +88,10 @@ mod _contextvars {
 
         fn enter(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
             if zelf.inner.entered.get() {
-                let msg = format!(
+                return Err(vm.new_runtime_error(format!(
                     "cannot enter context: {} is already entered",
                     zelf.as_object().repr(vm)?
-                );
-                return Err(vm.new_runtime_error(msg));
+                )));
             }
 
             super::CONTEXTS.with_borrow_mut(|ctxs| {
@@ -107,19 +105,20 @@ mod _contextvars {
 
         fn exit(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
             if !zelf.inner.entered.get() {
-                let msg = format!(
+                return Err(vm.new_runtime_error(format!(
                     "cannot exit context: {} is not entered",
                     zelf.as_object().repr(vm)?
-                );
-                return Err(vm.new_runtime_error(msg));
+                )));
             }
 
             super::CONTEXTS.with_borrow_mut(|ctxs| {
-                let err_msg =
-                    "cannot exit context: thread state references a different context object";
                 ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
                     .map(drop)
-                    .ok_or_else(|| vm.new_runtime_error(err_msg))
+                    .ok_or_else(|| {
+                        vm.new_runtime_error(
+                            "cannot exit context: thread state references a different context object"
+                        )
+                    })
             })?;
             zelf.inner.entered.set(false);
 
@@ -205,12 +204,11 @@ mod _contextvars {
             default: OptionalArg<PyObjectRef>,
         ) -> PyResult<Option<PyObjectRef>> {
             let found = self.get_inner(&key);
-            let result = if let Some(found) = found {
-                Some(found)
+            Ok(if found.is_some() {
+                found
             } else {
                 default.into_option()
-            };
-            Ok(result)
+            })
         }
 
         // TODO: wrong return type
@@ -252,12 +250,9 @@ mod _contextvars {
                 )),
                 subscript: atomic_func!(|mapping, needle, vm| {
                     let needle = needle.try_to_value(vm)?;
-                    let found = PyContext::mapping_downcast(mapping).get_inner(needle);
-                    if let Some(found) = found {
-                        Ok(found)
-                    } else {
-                        Err(vm.new_key_error(needle.to_owned().into()))
-                    }
+                    PyContext::mapping_downcast(mapping)
+                        .get_inner(needle)
+                        .ok_or_else(|| vm.new_key_error(needle.to_owned().into()))
                 }),
                 ass_subscript: None,
             };
@@ -333,8 +328,7 @@ mod _contextvars {
             if vars.swap_remove(zelf).is_none() {
                 // TODO:
                 // PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
-                let msg = zelf.as_object().repr(vm)?.as_wtf8().to_owned();
-                return Err(vm.new_lookup_error(msg));
+                return Err(vm.new_lookup_error(zelf.as_object().repr(vm)?.as_wtf8().to_owned()));
             }
 
             Ok(())
@@ -408,8 +402,7 @@ mod _contextvars {
             } else if let Some(default) = zelf.default.as_ref() {
                 default.clone()
             } else {
-                let msg = zelf.as_object().repr(vm)?;
-                return Err(vm.new_lookup_error(msg.as_wtf8().to_owned()));
+                return Err(vm.new_lookup_error(zelf.as_object().repr(vm)?.as_wtf8().to_owned()));
             };
             Ok(Some(value))
         }
@@ -439,25 +432,25 @@ mod _contextvars {
         #[pymethod]
         fn reset(zelf: &Py<Self>, token: PyRef<ContextToken>, vm: &VirtualMachine) -> PyResult<()> {
             if token.used.get() {
-                let msg = format!("{} has already been used once", token.as_object().repr(vm)?);
-                return Err(vm.new_runtime_error(msg));
+                return Err(vm.new_runtime_error(format!(
+                    "{} has already been used once",
+                    token.as_object().repr(vm)?
+                )));
             }
 
             if !zelf.is(&token.var) {
-                let msg = format!(
+                return Err(vm.new_value_error(format!(
                     "{} was created by a different ContextVar",
                     token.var.as_object().repr(vm)?
-                );
-                return Err(vm.new_value_error(msg));
+                )));
             }
 
             let ctx = PyContext::current(vm);
             if !ctx.is(&token.ctx) {
-                let msg = format!(
+                return Err(vm.new_value_error(format!(
                     "{} was created in a different Context",
                     token.var.as_object().repr(vm)?
-                );
-                return Err(vm.new_value_error(msg));
+                )));
             }
 
             token.used.set(true);
@@ -602,6 +595,7 @@ mod _contextvars {
         fn slot_new(_cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
             Err(vm.new_runtime_error("Tokens can only be created by ContextVars"))
         }
+
         fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
             unimplemented!("use slot_new")
         }
