diff --git a/src/main.rs b/src/main.rs index 5620d851e..3e77ddbc2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -147,9 +147,23 @@ fn test_run_script() { fn shell_exec(vm: &VirtualMachine, source: &str, scope: Scope) -> Result<(), CompileError> { match compile::compile(vm, source, &compile::Mode::Single, "".to_string()) { Ok(code) => { - if let Err(err) = vm.run_code_obj(code, scope) { - print_exception(vm, &err); + match vm.run_code_obj(code, scope.clone()) { + Ok(value) => { + // Save non-None values as "_" + + use rustpython_vm::pyobject::{IdProtocol, IntoPyObject, ItemProtocol}; + + if !value.is(&vm.get_none()) { + let key = objstr::PyString::from("_").into_pyobject(vm); + scope.globals.set_item(key, value, vm).unwrap(); + } + } + + Err(err) => { + print_exception(vm, &err); + } } + Ok(()) } // Don't inject syntax errors for line continuation diff --git a/vm/src/compile.rs b/vm/src/compile.rs index 76744f112..36b4c3ce7 100644 --- a/vm/src/compile.rs +++ b/vm/src/compile.rs @@ -130,18 +130,35 @@ impl Compiler { symbol_scope: SymbolScope, ) -> Result<(), CompileError> { self.scope_stack.push(symbol_scope); - for statement in &program.statements { + + let mut emitted_return = false; + + for (i, statement) in program.statements.iter().enumerate() { + let is_last = i == program.statements.len() - 1; + if let ast::Statement::Expression { ref expression } = statement.node { self.compile_expression(expression)?; - self.emit(Instruction::PrintExpr); + + if is_last { + self.emit(Instruction::Duplicate); + self.emit(Instruction::PrintExpr); + self.emit(Instruction::ReturnValue); + emitted_return = true; + } else { + self.emit(Instruction::PrintExpr); + } } else { self.compile_statement(&statement)?; } } - self.emit(Instruction::LoadConst { - value: bytecode::Constant::None, - }); - self.emit(Instruction::ReturnValue); + + if !emitted_return { + self.emit(Instruction::LoadConst { + value: bytecode::Constant::None, + }); + self.emit(Instruction::ReturnValue); + } + Ok(()) }