Skip to content

Bind interpreter queues with FromArgs #9729

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Nov 8, 2026
+80 −107
Merged

Bind interpreter queues with FromArgs #9729

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
187 changes: 80 additions & 107 deletions crates/vm/src/stdlib/_interpqueues.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,13 @@ pub(crate) use _interpqueues::{
#[pymodule]
pub(crate) mod _interpqueues {
use crate::{
AsObject, Py, PyObject, PyPayload, PyResult, VirtualMachine,
AsObject, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
builtins::{PyBaseExceptionRef, PyModule, PyType, PyTypeRef},
function::{ArgSpec, FuncArgs},
function::OptionalArg,
types::Constructor,
vm::crossinterp::{Fallback, SharedValue, UnboundOp},
};
use alloc::{collections::BTreeMap, collections::VecDeque, sync::Arc};
use core::cell::Cell;
use num_traits::{Signed, ToPrimitive};
use parking_lot::Mutex;
use std::sync::OnceLock;
Expand Down Expand Up @@ -337,23 +336,11 @@ pub(crate) mod _interpqueues {
Err(vm.new_overflow_error(format!("max queue ID is {}, got {repr}", i64::MAX)))
}

fn qid_arg(args: &FuncArgs, func: &'static str, vm: &VirtualMachine) -> PyResult<i64> {
let qid = Cell::new(None);
ArgSpec {
fname: func,
keywords: &["qid"],
required: 1,
max_positional: 1,
fn optional_int_arg(arg: &OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<i32> {
match arg {
OptionalArg::Present(obj) => int_arg(obj, vm),
OptionalArg::Missing => Ok(-1),
}
.parse_with(
args,
|_, obj, vm| {
qid.set(Some(parse_qid(obj, vm)?));
Ok(())
},
vm,
)?;
Ok(qid.get().unwrap())
}

#[expect(clippy::unnecessary_wraps, reason = "Needs to comply with a signature")]
Expand All @@ -363,39 +350,58 @@ pub(crate) mod _interpqueues {
Ok(())
}

// Keep the original inputs alive while applying the queue-specific converters.
#[derive(FromArgs)]
struct QueueIdArgs {
#[pyarg(any)]
qid: PyObjectRef,
}

#[derive(FromArgs)]
struct CreateArgs {
#[pyarg(any)]
maxsize: PyObjectRef,
#[pyarg(any, optional, py_default = "-1")]
unboundop: OptionalArg<PyObjectRef>,
#[pyarg(any, optional, py_default = "-1")]
fallback: OptionalArg<PyObjectRef>,
}

#[derive(FromArgs)]
struct PutArgs {
#[pyarg(any)]
qid: PyObjectRef,
#[pyarg(any)]
obj: PyObjectRef,
#[pyarg(any, optional, py_default = "-1")]
unboundop: OptionalArg<PyObjectRef>,
#[pyarg(any, optional, py_default = "-1")]
fallback: OptionalArg<PyObjectRef>,
}

#[derive(FromArgs)]
struct RegisterHeapTypesArgs {
#[pyarg(any)]
queuetype: PyObjectRef,
#[pyarg(any)]
emptyerror: PyObjectRef,
#[pyarg(any)]
fullerror: PyObjectRef,
}

#[pyfunction]
fn create(args: FuncArgs, vm: &VirtualMachine) -> PyResult<i64> {
let maxsize = Cell::new(None);
let unboundarg = Cell::new(None);
let fallbackarg = Cell::new(None);
ArgSpec {
fname: "create",
keywords: &["maxsize", "unboundop", "fallback"],
required: 1,
max_positional: 3,
}
.parse_with(
&args,
|i, obj, vm| {
match i {
0 => maxsize.set(Some(ssize_arg(obj, vm)?)),
1 => unboundarg.set(Some(int_arg(obj, vm)?)),
2 => fallbackarg.set(Some(int_arg(obj, vm)?)),
_ => unreachable!(),
}
Ok(())
},
vm,
)?;
let unboundop = resolve_unboundop(unboundarg.get().unwrap_or(-1), UnboundOp::Replace, vm)?;
let fallback =
resolve_fallback(fallbackarg.get().unwrap_or(-1), Fallback::Full.as_i32(), vm)?;
queue_create(maxsize.get().unwrap(), unboundop, fallback).map_err(|err| err.into_py(-1, vm))
fn create(args: CreateArgs, vm: &VirtualMachine) -> PyResult<i64> {
let maxsize = ssize_arg(&args.maxsize, vm)?;
let unboundarg = optional_int_arg(&args.unboundop, vm)?;
let fallbackarg = optional_int_arg(&args.fallback, vm)?;
let unboundop = resolve_unboundop(unboundarg, UnboundOp::Replace, vm)?;
let fallback = resolve_fallback(fallbackarg, Fallback::Full.as_i32(), vm)?;
queue_create(maxsize, unboundop, fallback).map_err(|err| err.into_py(-1, vm))
}

#[pyfunction]
fn destroy(args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = qid_arg(&args, "destroy", vm)?;
fn destroy(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = parse_qid(&args.qid, vm)?;
queue_destroy(qid).map_err(|err| err.into_py(qid, vm))
}

Expand Down Expand Up @@ -427,33 +433,10 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn put(args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = Cell::new(None);
let unboundarg = Cell::new(None);
let fallbackarg = Cell::new(None);
let parsed = ArgSpec {
fname: "put",
keywords: &["qid", "obj", "unboundop", "fallback"],
required: 2,
max_positional: 4,
}
.parse_with(
&args,
|i, obj, vm| {
match i {
0 => qid.set(Some(parse_qid(obj, vm)?)),
2 => unboundarg.set(Some(int_arg(obj, vm)?)),
3 => fallbackarg.set(Some(int_arg(obj, vm)?)),
_ => {}
}
Ok(())
},
vm,
)?;
let qid = qid.get().unwrap();
let obj = parsed[1].as_deref().unwrap();
let unboundarg = unboundarg.get().unwrap_or(-1);
let fallbackarg = fallbackarg.get().unwrap_or(-1);
fn put(args: PutArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = parse_qid(&args.qid, vm)?;
let unboundarg = optional_int_arg(&args.unboundop, vm)?;
let fallbackarg = optional_int_arg(&args.fallback, vm)?;
// The queue is only consulted when one of the arguments needs its default.
let (default_unboundop, default_fallback) = if unboundarg < 0 || fallbackarg < 0 {
let (unboundop, fallback) = queue_defaults(qid).map_err(|err| err.into_py(qid, vm))?;
Expand All @@ -463,7 +446,7 @@ pub(crate) mod _interpqueues {
};
let unboundop = resolve_unboundop(unboundarg, default_unboundop, vm)?;
let fallback = resolve_fallback(fallbackarg, default_fallback, vm)?;
let value = SharedValue::from_object(obj, fallback, vm)?;
let value = SharedValue::from_object(&args.obj, fallback, vm)?;
let queue = queue_lookup(qid).map_err(|err| err.into_py(qid, vm))?;
let mut state = queue.lock();
if !state.alive {
Expand All @@ -482,8 +465,8 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn get(args: FuncArgs, vm: &VirtualMachine) -> PyResult {
let qid = qid_arg(&args, "get", vm)?;
fn get(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult {
let qid = parse_qid(&args.qid, vm)?;
let item = {
let queue = queue_lookup(qid).map_err(|err| err.into_py(qid, vm))?;
let mut state = queue.lock();
Expand All @@ -509,20 +492,20 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn bind(args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = qid_arg(&args, "bind", vm)?;
fn bind(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = parse_qid(&args.qid, vm)?;
queue_bind(qid).map_err(|err| err.into_py(qid, vm))
}

#[pyfunction]
fn release(args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = qid_arg(&args, "release", vm)?;
fn release(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<()> {
let qid = parse_qid(&args.qid, vm)?;
queue_release(qid).map_err(|err| err.into_py(qid, vm))
}

#[pyfunction]
fn get_maxsize(args: FuncArgs, vm: &VirtualMachine) -> PyResult<isize> {
let qid = qid_arg(&args, "get_maxsize", vm)?;
fn get_maxsize(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<isize> {
let qid = parse_qid(&args.qid, vm)?;
let queue = queue_lookup(qid).map_err(|err| err.into_py(qid, vm))?;
let state = queue.lock();
if !state.alive {
Expand All @@ -532,8 +515,8 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn get_queue_defaults(args: FuncArgs, vm: &VirtualMachine) -> PyResult {
let qid = qid_arg(&args, "get_queue_defaults", vm)?;
fn get_queue_defaults(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult {
let qid = parse_qid(&args.qid, vm)?;
let (unboundop, fallback) = queue_defaults(qid).map_err(|err| err.into_py(qid, vm))?;
Ok(vm
.ctx
Expand All @@ -545,8 +528,8 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn is_full(args: FuncArgs, vm: &VirtualMachine) -> PyResult<bool> {
let qid = qid_arg(&args, "is_full", vm)?;
fn is_full(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<bool> {
let qid = parse_qid(&args.qid, vm)?;
let queue = queue_lookup(qid).map_err(|err| err.into_py(qid, vm))?;
let state = queue.lock();
if !state.alive {
Expand All @@ -556,8 +539,8 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn get_count(args: FuncArgs, vm: &VirtualMachine) -> PyResult<usize> {
let qid = qid_arg(&args, "get_count", vm)?;
fn get_count(args: QueueIdArgs, vm: &VirtualMachine) -> PyResult<usize> {
let qid = parse_qid(&args.qid, vm)?;
let queue = queue_lookup(qid).map_err(|err| err.into_py(qid, vm))?;
let state = queue.lock();
if !state.alive {
Expand All @@ -567,29 +550,19 @@ pub(crate) mod _interpqueues {
}

#[pyfunction]
fn _register_heap_types(args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
let parsed = ArgSpec {
fname: "_register_heap_types",
keywords: &["queuetype", "emptyerror", "fullerror"],
required: 3,
max_positional: 3,
}
.parse(&args, vm)?;
let queue = parsed[0]
.clone()
.unwrap()
fn _register_heap_types(args: RegisterHeapTypesArgs, vm: &VirtualMachine) -> PyResult<()> {
let queue = args
.queuetype
.downcast::<PyType>()
.map_err(|_| vm.new_type_error("expected a type for 'queuetype'"))?;
let empty = parsed[1]
.clone()
.unwrap()
let empty = args
.emptyerror
.downcast::<PyType>()
.ok()
.filter(|ty| ty.fast_issubclass(vm.ctx.exceptions.base_exception_type))
.ok_or_else(|| vm.new_type_error("expected an exception type for 'emptyerror'"))?;
let full = parsed[2]
.clone()
.unwrap()
let full = args
.fullerror
.downcast::<PyType>()
.ok()
.filter(|ty| ty.fast_issubclass(vm.ctx.exceptions.base_exception_type))
Expand Down
Loading