Skip to content

Commit 5c26a25

Browse files
coolreader18coolreader18
committed
Add async with, reorganize how with blocks work
1 parent bf08f2f commit 5c26a25

5 files changed

Lines changed: 132 additions & 115 deletions

File tree

‎bytecode/src/bytecode.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -218,9 +218,8 @@ pub enum Instruction {
218218
SetupWith {
219219
end: Label,
220220
},
221-
CleanupWith {
222-
end: Label,
223-
},
221+
WithCleanupStart,
222+
WithCleanupFinish,
224223
PopBlock,
225224
Raise {
226225
argc: usize,
@@ -275,6 +274,10 @@ pub enum Instruction {
275274
amount: usize,
276275
},
277276
GetAwaitable,
277+
BeforeAsyncWith,
278+
SetupAsyncWith {
279+
end: Label,
280+
},
278281
}
279282

280283
use self::Instruction::*;
@@ -528,7 +531,10 @@ impl Instruction {
528531
EnterFinally => w!(EnterFinally),
529532
EndFinally => w!(EndFinally),
530533
SetupWith { end } => w!(SetupWith, label_map[end]),
531-
CleanupWith { end } => w!(CleanupWith, label_map[end]),
534+
WithCleanupStart => w!(WithCleanupStart),
535+
WithCleanupFinish => w!(WithCleanupFinish),
536+
BeforeAsyncWith => w!(BeforeAsyncWith),
537+
SetupAsyncWith { end } => w!(SetupAsyncWith, label_map[end]),
532538
PopBlock => w!(PopBlock),
533539
Raise { argc } => w!(Raise, argc),
534540
BuildString { size } => w!(BuildString, size),

‎compiler/src/compile.rs‎

Lines changed: 61 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -386,27 +386,72 @@ impl<O: OutputStream> Compiler<O> {
386386
body,
387387
} => {
388388
if *is_async {
389-
unimplemented!("async with");
390-
} else {
391-
let end_label = self.new_label();
392-
for item in items {
393-
self.compile_expression(&item.context_expr)?;
394-
self.emit(Instruction::SetupWith { end: end_label });
395-
match &item.optional_vars {
396-
Some(var) => {
397-
self.compile_store(var)?;
398-
}
399-
None => {
400-
self.emit(Instruction::Pop);
389+
let end_labels = items
390+
.iter()
391+
.map(|item| {
392+
let end_label = self.new_label();
393+
self.compile_expression(&item.context_expr)?;
394+
self.emit(Instruction::BeforeAsyncWith);
395+
self.emit(Instruction::GetAwaitable);
396+
self.emit(Instruction::LoadConst {
397+
value: bytecode::Constant::None,
398+
});
399+
self.emit(Instruction::YieldFrom);
400+
self.emit(Instruction::SetupAsyncWith { end: end_label });
401+
match &item.optional_vars {
402+
Some(var) => {
403+
self.compile_store(var)?;
404+
}
405+
None => {
406+
self.emit(Instruction::Pop);
407+
}
401408
}
402-
}
409+
Ok(end_label)
410+
})
411+
.collect::<Result<Vec<_>, CompileError>>()?;
412+
413+
self.compile_statements(body)?;
414+
415+
for end_label in end_labels {
416+
self.emit(Instruction::PopBlock);
417+
self.emit(Instruction::EnterFinally);
418+
self.set_label(end_label);
419+
self.emit(Instruction::WithCleanupStart);
420+
self.emit(Instruction::GetAwaitable);
421+
self.emit(Instruction::LoadConst {
422+
value: bytecode::Constant::None,
423+
});
424+
self.emit(Instruction::YieldFrom);
425+
self.emit(Instruction::WithCleanupFinish);
403426
}
427+
} else {
428+
let end_labels = items
429+
.iter()
430+
.map(|item| {
431+
let end_label = self.new_label();
432+
self.compile_expression(&item.context_expr)?;
433+
self.emit(Instruction::SetupWith { end: end_label });
434+
match &item.optional_vars {
435+
Some(var) => {
436+
self.compile_store(var)?;
437+
}
438+
None => {
439+
self.emit(Instruction::Pop);
440+
}
441+
}
442+
Ok(end_label)
443+
})
444+
.collect::<Result<Vec<_>, CompileError>>()?;
404445

405446
self.compile_statements(body)?;
406-
for _ in 0..items.len() {
407-
self.emit(Instruction::CleanupWith { end: end_label });
447+
448+
for end_label in end_labels {
449+
self.emit(Instruction::PopBlock);
450+
self.emit(Instruction::EnterFinally);
451+
self.set_label(end_label);
452+
self.emit(Instruction::WithCleanupStart);
453+
self.emit(Instruction::WithCleanupFinish);
408454
}
409-
self.set_label(end_label);
410455
}
411456
}
412457
For {

‎vm/src/builtins.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -873,6 +873,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
873873
"FileNotFoundError" => ctx.exceptions.file_not_found_error.clone(),
874874
"FileExistsError" => ctx.exceptions.file_exists_error.clone(),
875875
"StopIteration" => ctx.exceptions.stop_iteration.clone(),
876+
"StopAsyncIteration" => ctx.exceptions.stop_async_iteration.clone(),
876877
"SystemError" => ctx.exceptions.system_error.clone(),
877878
"PermissionError" => ctx.exceptions.permission_error.clone(),
878879
"UnicodeError" => ctx.exceptions.unicode_error.clone(),

‎vm/src/exceptions.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,7 @@ pub struct ExceptionZoo {
244244
pub reference_error: PyClassRef,
245245
pub runtime_error: PyClassRef,
246246
pub stop_iteration: PyClassRef,
247+
pub stop_async_iteration: PyClassRef,
247248
pub syntax_error: PyClassRef,
248249
pub indentation_error: PyClassRef,
249250
pub tab_error: PyClassRef,
@@ -292,6 +293,7 @@ impl ExceptionZoo {
292293
let runtime_error = create_type("RuntimeError", &type_type, &exception_type);
293294
let reference_error = create_type("ReferenceError", &type_type, &exception_type);
294295
let stop_iteration = create_type("StopIteration", &type_type, &exception_type);
296+
let stop_async_iteration = create_type("StopAsyncIteration", &type_type, &exception_type);
295297
let syntax_error = create_type("SyntaxError", &type_type, &exception_type);
296298
let system_error = create_type("SystemError", &type_type, &exception_type);
297299
let type_error = create_type("TypeError", &type_type, &exception_type);
@@ -352,6 +354,7 @@ impl ExceptionZoo {
352354
permission_error,
353355
runtime_error,
354356
stop_iteration,
357+
stop_async_iteration,
355358
syntax_error,
356359
indentation_error,
357360
tab_error,

‎vm/src/frame.rs‎

Lines changed: 57 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,6 @@ enum BlockType {
4949
FinallyHandler {
5050
reason: Option<UnwindReason>,
5151
},
52-
With {
53-
end: bytecode::Label,
54-
context_manager: PyObjectRef,
55-
},
5652
ExceptHandler,
5753
}
5854

@@ -431,29 +427,69 @@ impl Frame {
431427
}
432428
bytecode::Instruction::SetupWith { end } => {
433429
let context_manager = self.pop_value();
430+
let exit = vm.get_attribute(context_manager.clone(), "__exit__")?;
431+
self.push_value(exit);
434432
// Call enter:
435-
let obj = vm.call_method(&context_manager, "__enter__", vec![])?;
436-
self.push_block(BlockType::With {
437-
end: *end,
438-
context_manager: context_manager.clone(),
433+
let enter_res = vm.call_method(&context_manager, "__enter__", vec![])?;
434+
self.push_block(BlockType::Finally { handler: *end });
435+
self.push_value(enter_res);
436+
Ok(None)
437+
}
438+
bytecode::Instruction::BeforeAsyncWith => {
439+
let mgr = self.pop_value();
440+
let aexit = vm.get_attribute(mgr.clone(), "__aexit__")?;
441+
self.push_value(aexit);
442+
let aenter_res = vm.call_method(&mgr, "__aenter__", vec![])?;
443+
self.push_value(aenter_res);
444+
445+
Ok(None)
446+
}
447+
bytecode::Instruction::SetupAsyncWith { end } => {
448+
self.push_block(BlockType::Finally { handler: *end });
449+
Ok(None)
450+
}
451+
bytecode::Instruction::WithCleanupStart => {
452+
let block = self.current_block().unwrap();
453+
let reason = match block.typ {
454+
BlockType::FinallyHandler { reason } => reason,
455+
_ => panic!("WithCleanupStart expects a FinallyHandler block on stack"),
456+
};
457+
let exc = reason.and_then(|reason| match reason {
458+
UnwindReason::Raising { exception } => Some(exception),
459+
_ => None,
439460
});
440-
self.push_value(obj);
461+
462+
let exit = self.pop_value();
463+
464+
let args = if let Some(exc) = exc {
465+
let exc_type = exc.class().into_object();
466+
let exc_val = exc.clone();
467+
let exc_tb = vm.ctx.none(); // TODO: retrieve traceback?
468+
vec![exc_type, exc_val, exc_tb]
469+
} else {
470+
vec![vm.ctx.none(), vm.ctx.none(), vm.ctx.none()]
471+
};
472+
let exit_res = vm.invoke(&exit, args)?;
473+
self.push_value(exit_res);
474+
441475
Ok(None)
442476
}
443-
bytecode::Instruction::CleanupWith { end: end1 } => {
477+
bytecode::Instruction::WithCleanupFinish => {
444478
let block = self.pop_block();
445-
if let BlockType::With {
446-
end: end2,
447-
context_manager,
448-
} = &block.typ
449-
{
450-
debug_assert!(end1 == end2);
451-
self.call_context_manager_exit(vm, &context_manager, None)?;
479+
let reason = match block.typ {
480+
BlockType::FinallyHandler { reason } => reason,
481+
_ => panic!("WithCleanupFinish expects a FinallyHandler block on stack"),
482+
};
483+
484+
let suppress_exception = objbool::boolval(vm, self.pop_value())?;
485+
if suppress_exception {
486+
// suppress exception
487+
Ok(None)
488+
} else if let Some(reason) = reason {
489+
self.unwind_blocks(vm, reason)
452490
} else {
453-
unreachable!("Block stack is incorrect, expected a with block");
491+
Ok(None)
454492
}
455-
456-
Ok(None)
457493
}
458494
bytecode::Instruction::PopBlock => {
459495
self.pop_block();
@@ -663,7 +699,7 @@ impl Frame {
663699
/// unwinding a block.
664700
/// Optionally returns an exception.
665701
#[cfg_attr(feature = "flame-it", flame("Frame"))]
666-
fn unwind_blocks(&self, vm: &VirtualMachine, mut reason: UnwindReason) -> FrameResult {
702+
fn unwind_blocks(&self, vm: &VirtualMachine, reason: UnwindReason) -> FrameResult {
667703
// First unwind all existing blocks on the block stack:
668704
while let Some(block) = self.current_block() {
669705
match block.typ {
@@ -699,60 +735,6 @@ impl Frame {
699735
return Ok(None);
700736
}
701737
}
702-
BlockType::With {
703-
context_manager,
704-
end,
705-
} => {
706-
self.pop_block();
707-
match &reason {
708-
UnwindReason::Raising { exception } => {
709-
match self.call_context_manager_exit(
710-
vm,
711-
&context_manager,
712-
Some(exception.clone()),
713-
) {
714-
Ok(exit_result_obj) => {
715-
match objbool::boolval(vm, exit_result_obj) {
716-
// If __exit__ method returned True, suppress the exception and continue execution.
717-
Ok(suppress_exception) => {
718-
if suppress_exception {
719-
self.jump(end);
720-
return Ok(None);
721-
} else {
722-
// go on with the stack unwinding.
723-
}
724-
}
725-
Err(exit_exc) => {
726-
reason = UnwindReason::Raising {
727-
exception: exit_exc,
728-
};
729-
// return Err(exit_exc);
730-
}
731-
}
732-
}
733-
Err(exit_exc) => {
734-
// TODO: what about original exception?
735-
reason = UnwindReason::Raising {
736-
exception: exit_exc,
737-
};
738-
// return Err(exit_exc);
739-
}
740-
}
741-
}
742-
_ => {
743-
match self.call_context_manager_exit(vm, &context_manager, None) {
744-
Ok(..) => {}
745-
Err(exit_exc) => {
746-
// __exit__ went wrong,
747-
reason = UnwindReason::Raising {
748-
exception: exit_exc,
749-
};
750-
// return Err(exc);
751-
}
752-
}
753-
}
754-
}
755-
}
756738
BlockType::FinallyHandler { .. } => {
757739
self.pop_block();
758740
}
@@ -773,26 +755,6 @@ impl Frame {
773755
}
774756
}
775757

776-
fn call_context_manager_exit(
777-
&self,
778-
vm: &VirtualMachine,
779-
context_manager: &PyObjectRef,
780-
exc: Option<PyObjectRef>,
781-
) -> PyResult {
782-
// TODO: do we want to put the exit call on the stack?
783-
// TODO: what happens when we got an error during execution of __exit__?
784-
let (exc_type, exc_val, exc_tb) = if let Some(exc) = exc {
785-
let exc_type = exc.class().into_object();
786-
let exc_val = exc.clone();
787-
let exc_tb = vm.ctx.none(); // TODO: retrieve traceback?
788-
(exc_type, exc_val, exc_tb)
789-
} else {
790-
(vm.ctx.none(), vm.ctx.none(), vm.ctx.none())
791-
};
792-
793-
vm.call_method(context_manager, "__exit__", vec![exc_type, exc_val, exc_tb])
794-
}
795-
796758
fn store_name(
797759
&self,
798760
vm: &VirtualMachine,

0 commit comments

Comments
 (0)