Visitar URL original
Fix yield-from to close sub-iterator on GeneratorExit · RustPython/RustPython@d3de5ea · GitHub
Skip to content

Commit d3de5ea

Browse files
youknowoneyouknowone
committed
Fix yield-from to close sub-iterator on GeneratorExit
When a generator doing yield-from receives GeneratorExit (e.g. from the destructor), close the sub-iterator via close() instead of throw(). This matches gen_close_iter() behavior: for generators/ coroutines call close(), for other iterators look for a close() method. Fixes unraisable exception in test_generators causing env_changed. Also skip flaky ProcessPoolForkWaitTest.test_first_exception on macOS.
1 parent c6c4d45 commit d3de5ea

6 files changed

Lines changed: 48 additions & 39 deletions

File tree

‎Lib/test/test_concurrent_futures/test_wait.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ class ProcessPoolForkWaitTest(ProcessPoolForkWaitTest): # TODO: RUSTPYTHON
205205
def test_first_completed(self): super().test_first_completed() # TODO: RUSTPYTHON
206206
@unittest.skipIf(sys.platform == 'linux', "TODO: RUSTPYTHON Fatal Python error: Segmentation fault")
207207
def test_first_completed_some_already_completed(self): super().test_first_completed_some_already_completed() # TODO: RUSTPYTHON
208-
@unittest.skipIf(sys.platform == 'linux', "TODO: RUSTPYTHON flaky")
208+
@unittest.skipIf(sys.platform != 'win32', "TODO: RUSTPYTHON flaky")
209209
def test_first_exception(self): super().test_first_exception() # TODO: RUSTPYTHON
210210
@unittest.skipIf(sys.platform == 'linux', "TODO: RUSTPYTHON flaky")
211211
def test_first_exception_one_already_failed(self): super().test_first_exception_one_already_failed() # TODO: RUSTPYTHON

‎Lib/test/test_yield_from.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -786,7 +786,6 @@ def outer():
786786
repr(value),
787787
])
788788

789-
@unittest.expectedFailure # TODO: RUSTPYTHON
790789
def test_throwing_GeneratorExit_into_subgen_that_returns(self):
791790
"""
792791
Test throwing GeneratorExit into a subgenerator that
@@ -817,7 +816,6 @@ def g():
817816
"Enter f",
818817
])
819818

820-
@unittest.expectedFailure # TODO: RUSTPYTHON
821819
def test_throwing_GeneratorExit_into_subgenerator_that_yields(self):
822820
"""
823821
Test throwing GeneratorExit into a subgenerator that
@@ -1135,7 +1133,6 @@ def outer():
11351133
self.assertIsNone(caught.exception.__context__)
11361134
self.assert_stop_iteration(g)
11371135

1138-
@unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: GeneratorExit() is not GeneratorExit()
11391136
def test_close_and_throw_raise_generator_exit(self):
11401137

11411138
yielded_first = object()
@@ -1524,7 +1521,6 @@ def outer():
15241521
self.assertIsNone(caught.exception.__context__)
15251522
self.assert_stop_iteration(g)
15261523

1527-
@unittest.expectedFailure # TODO: RUSTPYTHON
15281524
def test_close_and_throw_return(self):
15291525

15301526
yielded_first = object()

‎crates/vm/src/frame.rs‎

Lines changed: 44 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -635,39 +635,58 @@ impl ExecutingFrame<'_> {
635635
exc_tb: PyObjectRef,
636636
) -> PyResult<ExecutionResult> {
637637
if let Some(jen) = self.yield_from_target() {
638-
// borrow checker shenanigans - we only need to use exc_type/val/tb if the following
639-
// variable is Some
640-
let thrower = if let Some(coro) = self.builtin_coro(jen) {
641-
Some(Either::A(coro))
638+
// Check if the exception is GeneratorExit (type or instance).
639+
// For GeneratorExit, close the sub-iterator instead of throwing.
640+
let is_gen_exit = if let Some(typ) = exc_type.downcast_ref::<PyType>() {
641+
typ.fast_issubclass(vm.ctx.exceptions.generator_exit)
642642
} else {
643-
vm.get_attribute_opt(jen.to_owned(), "throw")?
644-
.map(Either::B)
643+
exc_type.fast_isinstance(vm.ctx.exceptions.generator_exit)
645644
};
646-
if let Some(thrower) = thrower {
647-
let ret = match thrower {
648-
Either::A(coro) => coro
649-
.throw(jen, exc_type, exc_val, exc_tb, vm)
650-
.to_pyresult(vm),
651-
Either::B(meth) => meth.call((exc_type, exc_val, exc_tb), vm),
645+
646+
if is_gen_exit {
647+
// gen_close_iter: close the sub-iterator
648+
let close_result = if let Some(coro) = self.builtin_coro(jen) {
649+
coro.close(jen, vm).map(|_| ())
650+
} else if let Some(close_meth) = vm.get_attribute_opt(jen.to_owned(), "close")? {
651+
close_meth.call((), vm).map(|_| ())
652+
} else {
653+
Ok(())
652654
};
653-
return ret.map(ExecutionResult::Yield).or_else(|err| {
654-
// This pushes Py_None to stack and restarts evalloop in exception mode.
655-
// Stack before throw: [receiver] (YIELD_VALUE already popped yielded value)
656-
// After pushing None: [receiver, None]
657-
// Exception handler will push exc: [receiver, None, exc]
658-
// CLEANUP_THROW expects: [sub_iter, last_sent_val, exc]
655+
if let Err(err) = close_result {
659656
self.push_value(vm.ctx.none());
660-
661-
// Set __context__ on the exception (_PyErr_ChainStackItem)
662657
vm.contextualize_exception(&err);
663-
664-
// Use unwind_blocks to let exception table route to CLEANUP_THROW
665-
match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) {
658+
return match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) {
666659
Ok(None) => self.run(vm),
667660
Ok(Some(result)) => Ok(result),
668661
Err(exception) => Err(exception),
669-
}
670-
});
662+
};
663+
}
664+
// Fall through to throw_here to raise GeneratorExit in the generator
665+
} else {
666+
// For non-GeneratorExit, delegate throw to sub-iterator
667+
let thrower = if let Some(coro) = self.builtin_coro(jen) {
668+
Some(Either::A(coro))
669+
} else {
670+
vm.get_attribute_opt(jen.to_owned(), "throw")?
671+
.map(Either::B)
672+
};
673+
if let Some(thrower) = thrower {
674+
let ret = match thrower {
675+
Either::A(coro) => coro
676+
.throw(jen, exc_type, exc_val, exc_tb, vm)
677+
.to_pyresult(vm),
678+
Either::B(meth) => meth.call((exc_type, exc_val, exc_tb), vm),
679+
};
680+
return ret.map(ExecutionResult::Yield).or_else(|err| {
681+
self.push_value(vm.ctx.none());
682+
vm.contextualize_exception(&err);
683+
match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) {
684+
Ok(None) => self.run(vm),
685+
Ok(Some(result)) => Ok(result),
686+
Err(exception) => Err(exception),
687+
}
688+
});
689+
}
671690
}
672691
}
673692
// throw_here: no delegate has throw method, or not in yield-from

‎crates/vm/src/protocol/callable.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -97,11 +97,7 @@ impl core::fmt::Display for TraceEvent {
9797
impl VirtualMachine {
9898
/// Call registered trace function.
9999
#[inline]
100-
pub(crate) fn trace_event(
101-
&self,
102-
event: TraceEvent,
103-
arg: Option<PyObjectRef>,
104-
) -> PyResult<()> {
100+
pub(crate) fn trace_event(&self, event: TraceEvent, arg: Option<PyObjectRef>) -> PyResult<()> {
105101
if self.use_tracing.get() {
106102
self._trace_event_inner(event, arg)
107103
} else {

‎crates/vm/src/stdlib/io.rs‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5267,9 +5267,7 @@ mod fileio {
52675267
match fd {
52685268
Ok(fd) => (fd.into_raw(), Some(filename)),
52695269
Err(e) => {
5270-
return Err(OSErrorBuilder::with_filename_from_errno(
5271-
&e, filename, vm,
5272-
))
5270+
return Err(OSErrorBuilder::with_filename_from_errno(&e, filename, vm));
52735271
}
52745272
}
52755273
}

‎crates/vm/src/stdlib/itertools.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@ mod decl {
1616
stdlib::sys,
1717
types::{Constructor, IterNext, Iterable, Representable, SelfIter},
1818
};
19+
use core::sync::atomic::{AtomicBool, Ordering};
1920
use crossbeam_utils::atomic::AtomicCell;
2021
use malachite_bigint::BigInt;
21-
use core::sync::atomic::{AtomicBool, Ordering};
2222
use num_traits::One;
2323

2424
use alloc::fmt;

0 commit comments

Comments
 (0)