Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
iterable context
  • Loading branch information
youknowone committed Jan 14, 2026
commit 6ff5979e4af3fc6d4b2704dff933cad314ea7eab
45 changes: 37 additions & 8 deletions crates/stdlib/src/contextvars.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -163,7 +163,7 @@ mod _contextvars {
}
}

#[pyclass(with(Constructor, AsMapping, AsSequence))]
#[pyclass(with(Constructor, AsMapping, AsSequence, Iterable))]
impl PyContext {
#[pymethod]
fn run(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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 {
Expand Down