use crate::{
AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
builtins::{PyCode, PyStrRef, PyTraceback, PyTupleRef},
common::lock::PyMutex,
exceptions::types::PyBaseException,
frame::{ExecutionResult, FrameObject, FrameObjectRef, FrameOwner, InterpreterFrame},
function::OptionalArg,
object::{PyAtomicRef, Traverse, TraverseFn},
protocol::PyIterReturn,
vm::GenFrameLink,
};
use crossbeam_utils::atomic::AtomicCell;
impl ExecutionResult {
fn into_iter_return(self, vm: &VirtualMachine) -> PyIterReturn {
match self {
Self::Yield(value) => PyIterReturn::Return(value),
Self::Return(value) => {
let arg = if vm.is_none(&value) {
None
} else {
Some(value)
};
PyIterReturn::StopIteration(arg)
}
Self::TailCall | Self::GenResume => {
unreachable!("unfinished frame result in generator/coroutine")
}
}
}
}
#[derive(Debug)]
pub struct Coro {
frame: PyAtomicRef<Option<FrameObject>>,
code: PyRef<PyCode>,
pub closed: AtomicCell<bool>, running: AtomicCell<bool>,
name: PyMutex<PyStrRef>,
qualname: PyMutex<PyStrRef>,
exception: PyAtomicRef<Option<PyBaseException>>, }
unsafe impl Traverse for Coro {
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
if let Some(frame) = self.frame.deref() {
tracer_fn(frame.as_object());
}
self.code.traverse(tracer_fn);
self.name.traverse(tracer_fn);
self.qualname.traverse(tracer_fn);
if let Some(exc) = self.exception.deref() {
exc.traverse(tracer_fn);
}
}
}
struct RunningGuard<'a>(&'a Coro);
impl Drop for RunningGuard<'_> {
fn drop(&mut self) {
self.0.running.store(false);
}
}
#[inline]
pub(crate) fn as_builtin_coro(obj: &PyObject) -> Option<&Coro> {
use crate::builtins::{PyCoroutine, PyGenerator};
crate::match_class!(match obj {
ref g @ PyGenerator => Some(g.as_coro()),
ref c @ PyCoroutine => Some(c.as_coro()),
_ => None,
})
}
pub(crate) struct FlatResume {
jen: PyObjectRef,
frame: *const Py<FrameObject>,
link: GenFrameLink,
}
pub(crate) enum FlatEnter {
Entered {
state: FlatResume,
value: Option<PyObjectRef>,
},
Exhausted,
}
impl FlatResume {
pub(crate) fn iframe_ptr(&self) -> *mut InterpreterFrame {
unsafe { (*self.frame).iframe_mut() as *mut InterpreterFrame }
}
}
#[inline]
pub(crate) fn flat_resume_enter(
jen: PyObjectRef,
value: PyObjectRef,
vm: &VirtualMachine,
) -> PyResult<FlatEnter> {
let coro = as_builtin_coro(&jen).expect("flat resume of a non-generator");
if coro.closed.load() {
return Ok(FlatEnter::Exhausted);
}
let _claim = coro.claim(&jen, vm)?;
if coro.closed.load() {
return Ok(FlatEnter::Exhausted);
}
let value = if coro.frame_opt().is_some_and(|f| f.lasti() > 0) {
Some(value)
} else if !vm.is_none(&value) {
return Err(vm.new_type_error(format!(
"can't send non-None value to a just-started {}",
gen_name(&jen, vm),
)));
} else {
None
};
let gen_exc = unsafe { coro.exception.swap(None) };
let frame_ref = coro.frame();
let frame: *const Py<FrameObject> = &*frame_ref;
let link = match vm.gen_frame_link(&frame_ref, gen_exc) {
Ok(link) => link,
Err(exc) => {
coro.retire(&_claim);
return Err(exc);
}
};
core::mem::forget(_claim);
Ok(FlatEnter::Entered {
state: FlatResume { jen, frame, link },
value,
})
}
#[inline]
pub(crate) fn flat_resume_exit(
state: FlatResume,
result: PyResult<ExecutionResult>,
vm: &VirtualMachine,
) -> PyResult<PyIterReturn> {
let FlatResume { jen, frame, link } = state;
let coro = as_builtin_coro(&jen).expect("flat resume of a non-generator");
let _old = unsafe { coro.exception.swap(vm.current_exception()) };
vm.gen_frame_unlink(unsafe { &*frame }, link);
let claim = RunningGuard(coro);
coro.maybe_close(&result, &claim);
drop(claim);
coro.finalize_send_result(result, &jen, vm)
}
fn gen_name(jen: &PyObject, vm: &VirtualMachine) -> &'static str {
let typ = jen.class();
if typ.is(vm.ctx.types.coroutine_type) {
"coroutine"
} else if typ.is(vm.ctx.types.async_generator) {
"async generator"
} else {
"generator"
}
}
impl Coro {
pub fn new(frame: FrameObjectRef, name: PyStrRef, qualname: PyStrRef) -> Self {
let code = frame.iframe().code().to_owned();
frame.as_object().mark_cache_published();
Self {
frame: Some(frame).into(),
code,
closed: AtomicCell::new(false),
running: AtomicCell::new(false),
exception: PyAtomicRef::from(None),
name: PyMutex::new(name),
qualname: PyMutex::new(qualname),
}
}
fn clear_except_code(&self) {
let Some(frame) = (unsafe { self.frame.swap(None) }) else {
return;
};
frame.clear_generator();
if frame.as_object().strong_count() == 1 {
frame.clear_locals_and_stack();
} else {
frame.iframe().owner.store(
FrameOwner::FrameObject as i8,
core::sync::atomic::Ordering::Release,
);
}
}
fn maybe_close(&self, res: &PyResult<ExecutionResult>, _claim: &RunningGuard<'_>) {
match res {
Ok(ExecutionResult::Return(_)) | Err(_) => {
self.closed.store(true);
self.clear_except_code();
}
Ok(ExecutionResult::Yield(_)) => {}
Ok(ExecutionResult::TailCall | ExecutionResult::GenResume) => {
unreachable!("unfinished frame result in generator/coroutine")
}
}
}
fn retire(&self, _claim: &RunningGuard<'_>) {
self.closed.store(true);
self.clear_except_code();
}
fn claim(&self, jen: &PyObject, vm: &VirtualMachine) -> PyResult<RunningGuard<'_>> {
if self.running.compare_exchange(false, true).is_err() {
return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm))));
}
Ok(RunningGuard(self))
}
fn run_claimed<F>(
&self,
_claim: &RunningGuard<'_>,
vm: &VirtualMachine,
func: F,
) -> PyResult<ExecutionResult>
where
F: FnOnce(&Py<FrameObject>) -> PyResult<ExecutionResult>,
{
let gen_exc = unsafe { self.exception.swap(None) };
let exception_ptr = &self.exception as *const PyAtomicRef<Option<PyBaseException>>;
let frame = self.frame();
vm.resume_gen_frame(&frame, gen_exc, |f| {
let result = func(f);
let _old = unsafe { (*exception_ptr).swap(vm.current_exception()) };
result
})
}
fn finalize_send_result(
&self,
result: PyResult<ExecutionResult>,
jen: &PyObject,
vm: &VirtualMachine,
) -> PyResult<PyIterReturn> {
match result {
Ok(exec_res) => Ok(exec_res.into_iter_return(vm)),
Err(e) => {
if e.fast_isinstance(vm.ctx.exceptions.stop_iteration) {
let err =
vm.new_runtime_error(format!("{} raised StopIteration", gen_name(jen, vm)));
err.set_context(Some(e.clone()));
err.set_cause(Some(e));
Err(err)
} else if jen.class().is(vm.ctx.types.async_generator)
&& e.fast_isinstance(vm.ctx.exceptions.stop_async_iteration)
{
let err = vm.new_runtime_error("async generator raised StopAsyncIteration");
err.set_context(Some(e.clone()));
err.set_cause(Some(e));
Err(err)
} else {
Err(e)
}
}
}
}
pub(crate) fn send_none(&self, jen: &PyObject, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
if self.closed.load() {
return Self::send_when_closed(jen, vm);
}
let claim = self.claim(jen, vm)?;
if self.closed.load() {
return Self::send_when_closed(jen, vm);
}
let value = if self.frame_opt().is_some_and(|f| f.lasti() > 0) {
Some(vm.ctx.none())
} else {
None
};
let result = self.run_claimed(&claim, vm, |f| f.resume(value, vm));
self.maybe_close(&result, &claim);
drop(claim);
self.finalize_send_result(result, jen, vm)
}
fn send_when_closed(jen: &PyObject, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
if jen.class().is(vm.ctx.types.coroutine_type) {
Err(vm.new_runtime_error("cannot reuse already awaited coroutine"))
} else {
Ok(PyIterReturn::StopIteration(None))
}
}
fn throw_when_closed(
jen: &PyObject,
exc_type: PyObjectRef,
exc_val: PyObjectRef,
exc_tb: PyObjectRef,
vm: &VirtualMachine,
) -> PyResult<PyIterReturn> {
if jen.class().is(vm.ctx.types.coroutine_type) {
Err(vm.new_runtime_error("cannot reuse already awaited coroutine"))
} else {
Err(vm.normalize_exception(exc_type, exc_val, exc_tb)?)
}
}
pub fn send(
&self,
jen: &PyObject,
value: PyObjectRef,
vm: &VirtualMachine,
) -> PyResult<PyIterReturn> {
if self.closed.load() {
return Self::send_when_closed(jen, vm);
}
let claim = self.claim(jen, vm)?;
if self.closed.load() {
return Self::send_when_closed(jen, vm);
}
let value = if self.frame_opt().is_some_and(|f| f.lasti() > 0) {
Some(value)
} else if !vm.is_none(&value) {
return Err(vm.new_type_error(format!(
"can't send non-None value to a just-started {}",
gen_name(jen, vm),
)));
} else {
None
};
let result = self.run_claimed(&claim, vm, |f| f.resume(value, vm));
self.maybe_close(&result, &claim);
drop(claim);
self.finalize_send_result(result, jen, vm)
}
pub fn throw(
&self,
jen: &PyObject,
exc_type: PyObjectRef,
exc_val: PyObjectRef,
exc_tb: PyObjectRef,
vm: &VirtualMachine,
) -> PyResult<PyIterReturn> {
if exc_type.fast_isinstance(vm.ctx.exceptions.base_exception_type) && !vm.is_none(&exc_val)
{
return Err(vm.new_type_error("instance exception may not have a separate value"));
}
if !vm.is_none(&exc_tb) && !exc_tb.fast_isinstance(vm.ctx.types.traceback_type) {
return Err(vm.new_type_error("throw() third argument must be a traceback object"));
}
if self.closed.load() {
return Self::throw_when_closed(jen, exc_type, exc_val, exc_tb, vm);
}
crate::exceptions::ExceptionCtor::try_from_object(vm, exc_type.clone())?;
let claim = self.claim(jen, vm)?;
if self.closed.load() {
drop(claim);
return Self::throw_when_closed(jen, exc_type, exc_val, exc_tb, vm);
}
let result = self.run_claimed(&claim, vm, |f| f.gen_throw(vm, exc_type, exc_val, exc_tb));
self.maybe_close(&result, &claim);
drop(claim);
self.finalize_send_result(result, jen, vm)
}
pub fn close(&self, jen: &PyObject, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
if self.closed.load() {
return Ok(vm.ctx.none());
}
let claim = self.claim(jen, vm)?;
if self.closed.load() {
return Ok(vm.ctx.none());
}
if self.frame_opt().is_none_or(|f| f.lasti() == 0) {
self.closed.store(true);
self.clear_except_code();
return Ok(vm.ctx.none());
}
let result = self.run_claimed(&claim, vm, |f| {
f.gen_throw(
vm,
vm.ctx.exceptions.generator_exit.to_owned().into(),
vm.ctx.none(),
vm.ctx.none(),
)
});
if !matches!(&result, Ok(ExecutionResult::Yield(_))) {
self.closed.store(true);
self.clear_except_code();
}
drop(claim);
match result {
Ok(ExecutionResult::Yield(_)) => {
Err(vm.new_runtime_error(format!("{} ignored GeneratorExit", gen_name(jen, vm))))
}
Err(e) if !is_gen_exit(&e, vm) => Err(e),
Ok(ExecutionResult::Return(value)) => Ok(value),
_ => Ok(vm.ctx.none()),
}
}
pub(crate) fn started(&self) -> bool {
self.frame_opt().is_some_and(|f| f.lasti() > 0)
}
pub fn suspended(&self) -> bool {
!self.closed.load()
&& !self.running.load()
&& self.frame_opt().is_some_and(|f| f.lasti() > 0)
}
pub fn running(&self) -> bool {
self.running.load()
}
pub fn closed(&self) -> bool {
self.closed.load()
}
pub(crate) fn frame_ref(&self) -> &Py<FrameObject> {
self.frame.deref().expect("generator frame")
}
pub fn frame(&self) -> FrameObjectRef {
self.frame_opt().expect("generator frame")
}
pub fn frame_opt(&self) -> Option<FrameObjectRef> {
self.frame.load_owned()
}
pub fn code(&self) -> PyRef<PyCode> {
self.code.clone()
}
pub fn name(&self) -> PyStrRef {
self.name.lock().clone()
}
pub fn set_name(&self, name: PyStrRef) {
*self.name.lock() = name;
}
pub fn qualname(&self) -> PyStrRef {
self.qualname.lock().clone()
}
pub fn set_qualname(&self, qualname: PyStrRef) {
*self.qualname.lock() = qualname;
}
pub fn repr(&self, jen: &PyObject, id: usize, vm: &VirtualMachine) -> String {
let qualname = self.qualname();
format!(
"<{} object {} at {:#x}>",
gen_name(jen, vm),
qualname.as_wtf8(),
id
)
}
}
pub(crate) fn is_gen_exit(exc: &Py<PyBaseException>, vm: &VirtualMachine) -> bool {
exc.fast_isinstance(vm.ctx.exceptions.generator_exit)
}
fn iframe_origin_lineno(frame: &InterpreterFrame) -> usize {
let lasti = frame.get_lasti();
if lasti == 0 {
return frame.code().first_line_number.map_or(1, |n| n.get());
}
frame.code().locations[lasti as usize - 1].0.line.get()
}
pub(crate) fn compute_cr_origin(vm: &VirtualMachine) -> Option<PyTupleRef> {
let depth = crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.get();
if depth == 0 {
return None;
}
let mut items = Vec::new();
let mut iframe = crate::frame::current_thread_iframe();
while !iframe.is_null() && items.len() < depth as usize {
let frame = unsafe { &*iframe };
let code = frame.code();
items.push(
vm.ctx
.new_tuple(vec![
code.source_path().to_owned().into(),
vm.ctx.new_int(iframe_origin_lineno(frame)).into(),
code.obj_name.to_owned().into(),
])
.into(),
);
iframe = frame.previous();
}
if items.is_empty() {
None
} else {
Some(vm.ctx.new_tuple(items))
}
}
pub(crate) fn get_awaitable_iter(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
use crate::builtins::{PyCoroutine, PyGenerator};
use crate::protocol::PyIter;
if obj.downcastable::<PyCoroutine>()
|| obj.downcast_ref::<PyGenerator>().is_some_and(|g| {
g.as_coro()
.code()
.flags
.contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
})
{
return Ok(obj);
}
if let Some(await_method) = vm.get_method(obj.clone(), identifier!(vm, __await__)) {
let result = await_method?.call((), vm)?;
if result.downcastable::<PyCoroutine>()
|| result.downcast_ref::<PyGenerator>().is_some_and(|g| {
g.as_coro()
.code()
.flags
.contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
})
{
return Err(vm.new_type_error("__await__() returned a coroutine"));
}
if !PyIter::check(&result) {
return Err(vm.new_type_error(format!(
"__await__() returned non-iterator of type '{}'",
result.class().name()
)));
}
return Ok(result);
}
Err(vm.new_type_error(format!("'{}' object can't be awaited", obj.class().name())))
}
pub(crate) fn unraisable_while_closing(
jen: &PyObject,
coro: &Coro,
e: crate::builtins::PyBaseExceptionRef,
vm: &VirtualMachine,
) {
if e.traceback().is_none()
&& let Some(frame) = coro.frame_opt()
{
let lasti = frame.lasti().saturating_mul(2) as i32;
let lineno = rustpython_compiler_core::OneIndexed::new(frame.lineno().max(1) as usize)
.unwrap_or(rustpython_compiler_core::OneIndexed::MIN);
let tb = PyTraceback::new(None, frame, lasti, lineno);
e.set_traceback(Some(tb.into_ref(&vm.ctx)));
}
let msg = jen
.repr(vm)
.ok()
.map(|r| format!("Exception ignored while closing generator {r}"));
vm.run_unraisable(e, msg, vm.ctx.none());
}
pub(crate) fn warn_deprecated_throw_signature(
exc_val: &OptionalArg,
exc_tb: &OptionalArg,
vm: &VirtualMachine,
) -> PyResult<()> {
if exc_val.is_present() || exc_tb.is_present() {
crate::warn::warn(
vm.ctx
.new_str(
"the (type, val, tb) signature of throw() is deprecated, \
use throw(val) instead",
)
.into(),
Some(vm.ctx.exceptions.deprecation_warning.to_owned()),
1,
None,
vm,
)?;
}
Ok(())
}