Make PyResult<T> = Result<T, PyBaseExceptionRef>

Author: coolreader18

src/main.rs

@@ -7,7 +7,7 @@ extern crate log;
 use clap::{App, AppSettings, Arg, ArgMatches};
 use rustpython_compiler::compile;
 use rustpython_vm::{
-    exceptions::{print_exception, PyBaseExceptionRef},
+    exceptions::print_exception,
     match_class,
     obj::{objint::PyInt, objtype},
     pyobject::{ItemProtocol, PyResult},
@@ -51,7 +51,6 @@ fn main() {
     // See if any exception leaked out:
     if let Err(err) = res {
         if objtype::isinstance(&err, &vm.ctx.exceptions.system_exit) {
-            let err: PyBaseExceptionRef = err.downcast().unwrap();
             let args = err.args();
             match args.elements.len() {
                 0 => return,

src/shell.rs

@@ -5,9 +5,9 @@ mod rustyline_helper;
 use rustpython_compiler::{compile, error::CompileError, error::CompileErrorType};
 use rustpython_parser::error::ParseErrorType;
 use rustpython_vm::{
-    exceptions::print_exception,
+    exceptions::{print_exception, PyBaseExceptionRef},
     obj::objtype,
-    pyobject::{ItemProtocol, PyObjectRef, PyResult},
+    pyobject::{ItemProtocol, PyResult},
     scope::Scope,
     VirtualMachine,
 };
@@ -16,7 +16,7 @@ use readline::{Readline, ReadlineResult};
 
 enum ShellExecResult {
     Ok,
-    PyErr(PyObjectRef),
+    PyErr(PyBaseExceptionRef),
     Continue,
 }
 

vm/src/builtins.rs

@@ -12,6 +12,7 @@ use num_traits::{Signed, ToPrimitive, Zero};
 #[cfg(feature = "rustpython-compiler")]
 use rustpython_compiler::compile;
 
+use crate::exceptions::PyBaseExceptionRef;
 use crate::function::{single_or_tuple_any, Args, KwArgs, OptionalArg, PyFuncArgs};
 use crate::obj::objbool::{self, IntoPyBool};
 use crate::obj::objbyteinner::PyByteInner;
@@ -290,7 +291,7 @@ fn builtin_format(
         })
 }
 
-fn catch_attr_exception<T>(ex: PyObjectRef, default: T, vm: &VirtualMachine) -> PyResult<T> {
+fn catch_attr_exception<T>(ex: PyBaseExceptionRef, default: T, vm: &VirtualMachine) -> PyResult<T> {
     if objtype::isinstance(&ex, &vm.ctx.exceptions.attribute_error) {
         Ok(default)
     } else {

vm/src/exceptions.rs

@@ -4,8 +4,7 @@ use crate::obj::objtraceback::PyTracebackRef;
 use crate::obj::objtuple::{PyTuple, PyTupleRef};
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{
-    IdProtocol, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
-    TypeProtocol,
+    PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
 };
 use crate::types::create_type;
 use crate::vm::VirtualMachine;
@@ -18,8 +17,8 @@ use std::io::{self, BufRead, BufReader, Write};
 #[pyclass]
 pub struct PyBaseException {
     traceback: RefCell<Option<PyTracebackRef>>,
-    cause: RefCell<Option<PyObjectRef>>,
-    context: RefCell<Option<PyObjectRef>>,
+    cause: RefCell<Option<PyBaseExceptionRef>>,
+    context: RefCell<Option<PyBaseExceptionRef>>,
     suppress_context: Cell<bool>,
     args: RefCell<PyTupleRef>,
 }
@@ -83,12 +82,12 @@ impl PyBaseException {
     }
 
     #[pyproperty(name = "__traceback__", setter)]
-    fn set_traceback(&self, traceback: Option<PyTracebackRef>, _vm: &VirtualMachine) {
+    fn setter_traceback(&self, traceback: Option<PyTracebackRef>, _vm: &VirtualMachine) {
         self.traceback.replace(traceback);
     }
 
     #[pyproperty(name = "__cause__")]
-    fn get_cause(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
+    fn get_cause(&self, _vm: &VirtualMachine) -> Option<PyBaseExceptionRef> {
         self.cause.borrow().clone()
     }
 
@@ -98,7 +97,7 @@ impl PyBaseException {
     }
 
     #[pyproperty(name = "__context__")]
-    fn get_context(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
+    fn get_context(&self, _vm: &VirtualMachine) -> Option<PyBaseExceptionRef> {
         self.context.borrow().clone()
     }
 
@@ -150,51 +149,67 @@ impl PyBaseException {
     pub fn args(&self) -> PyTupleRef {
         self.args.borrow().clone()
     }
+
+    pub fn traceback(&self) -> Option<PyTracebackRef> {
+        self.traceback.borrow().clone()
+    }
+    pub fn set_traceback(&self, tb: Option<PyTracebackRef>) {
+        self.traceback.replace(tb);
+    }
+
+    pub fn cause(&self) -> Option<PyBaseExceptionRef> {
+        self.cause.borrow().clone()
+    }
+    pub fn set_cause(&self, cause: Option<PyBaseExceptionRef>) {
+        self.cause.replace(cause);
+    }
+
+    pub fn context(&self) -> Option<PyBaseExceptionRef> {
+        self.context.borrow().clone()
+    }
+    pub fn set_context(&self, context: Option<PyBaseExceptionRef>) {
+        self.context.replace(context);
+    }
 }
 
 /// Print exception chain
-pub fn print_exception(vm: &VirtualMachine, exc: &PyObjectRef) {
-    let _ = write_exception(io::stdout(), vm, exc);
+pub fn print_exception(vm: &VirtualMachine, exc: &PyBaseExceptionRef) {
+    let stdout = io::stdout();
+    let mut stdout = stdout.lock();
+    let _ = write_exception(&mut stdout, vm, exc);
 }
 
 pub fn write_exception<W: Write>(
-    mut output: W,
+    output: &mut W,
     vm: &VirtualMachine,
-    exc: &PyObjectRef,
+    exc: &PyBaseExceptionRef,
 ) -> io::Result<()> {
-    let mut had_cause = false;
-    if let Ok(cause) = vm.get_attribute(exc.clone(), "__cause__") {
-        if !vm.get_none().is(&cause) {
-            had_cause = true;
-            print_exception(vm, &cause);
+    let cause = exc.cause();
+    if let Some(cause) = &cause {
+        write_exception(output, vm, cause)?;
+        writeln!(
+            output,
+            "\nThe above exception was the direct cause of the following exception:\n"
+        )?;
+    }
+    if cause.is_none() {
+        if let Some(context) = exc.context() {
+            write_exception(output, vm, &context)?;
             writeln!(
                 output,
-                "\nThe above exception was the direct cause of the following exception:\n"
+                "\nDuring handling of the above exception, another exception occurred:\n"
             )?;
         }
     }
-    if !had_cause {
-        if let Ok(context) = vm.get_attribute(exc.clone(), "__context__") {
-            if !vm.get_none().is(&context) {
-                print_exception(vm, &context);
-                writeln!(
-                    output,
-                    "\nDuring handling of the above exception, another exception occurred:\n"
-                )?;
-            }
-        }
-    }
-    print_exception_inner(output, vm, exc)
+    write_exception_inner(output, vm, exc)
 }
 
-fn print_source_line<W: Write>(mut output: W, filename: &str, lineno: usize) -> io::Result<()> {
+fn print_source_line<W: Write>(output: &mut W, filename: &str, lineno: usize) -> io::Result<()> {
     // TODO: use io.open() method instead, when available, according to https://github.com/python/cpython/blob/master/Python/traceback.c#L393
     // TODO: support different encodings
     let file = match File::open(filename) {
         Ok(file) => file,
-        Err(_) => {
-            return Ok(());
-        }
+        Err(_) => return Ok(()),
     };
     let file = BufReader::new(file);
 
@@ -212,7 +227,7 @@ fn print_source_line<W: Write>(mut output: W, filename: &str, lineno: usize) ->
 }
 
 /// Print exception occurrence location from traceback element
-fn print_traceback_entry<W: Write>(mut output: W, tb_entry: &PyTracebackRef) -> io::Result<()> {
+fn write_traceback_entry<W: Write>(output: &mut W, tb_entry: &PyTracebackRef) -> io::Result<()> {
     let filename = tb_entry.frame.code.source_path.to_string();
     writeln!(
         output,
@@ -225,19 +240,17 @@ fn print_traceback_entry<W: Write>(mut output: W, tb_entry: &PyTracebackRef) ->
 }
 
 /// Print exception with traceback
-pub fn print_exception_inner<W: Write>(
-    mut output: W,
+pub fn write_exception_inner<W: Write>(
+    output: &mut W,
     vm: &VirtualMachine,
-    exc: &PyObjectRef,
+    exc: &PyBaseExceptionRef,
 ) -> io::Result<()> {
-    let exc: PyBaseExceptionRef = exc.clone().downcast().unwrap();
-
     if let Some(tb) = exc.traceback.borrow().clone() {
         writeln!(output, "Traceback (most recent call last):")?;
-        let mut tb = &Some(tb);
+        let mut tb = Some(&tb);
         while let Some(traceback) = tb {
-            print_traceback_entry(&mut output, traceback)?;
-            tb = &traceback.next;
+            write_traceback_entry(output, traceback)?;
+            tb = traceback.next.as_ref();
         }
     } else {
         writeln!(output, "No traceback set on exception")?;
@@ -505,14 +518,14 @@ fn none_getter(_obj: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
     vm.get_none()
 }
 
-fn make_arg_getter(idx: usize) -> impl Fn(PyBaseExceptionRef, &VirtualMachine) -> PyResult {
+fn make_arg_getter(idx: usize) -> impl Fn(PyBaseExceptionRef, &VirtualMachine) -> PyObjectRef {
     move |exc, vm| {
         exc.args
             .borrow()
             .elements
             .get(idx)
             .cloned()
-            .ok_or_else(|| vm.new_value_error(format!("couldn't get arg {} of exception", idx)))
+            .unwrap_or_else(|| vm.get_none())
     }
 }
 

vm/src/frame.rs

@@ -5,6 +5,7 @@ use indexmap::IndexMap;
 use itertools::Itertools;
 
 use crate::bytecode;
+use crate::exceptions::PyBaseExceptionRef;
 use crate::function::{single_or_tuple_any, PyFuncArgs};
 use crate::obj::objbool;
 use crate::obj::objcode::PyCodeRef;
@@ -15,7 +16,7 @@ use crate::obj::objiter;
 use crate::obj::objlist;
 use crate::obj::objslice::PySlice;
 use crate::obj::objstr::{self, PyString};
-use crate::obj::objtraceback::{PyTraceback, PyTracebackRef};
+use crate::obj::objtraceback::PyTraceback;
 use crate::obj::objtuple::PyTuple;
 use crate::obj::objtype::{self, PyClassRef};
 use crate::pyobject::{
@@ -63,7 +64,7 @@ enum UnwindReason {
     Returning { value: PyObjectRef },
 
     /// We hit an exception, so unwind any try-except and finally blocks.
-    Raising { exception: PyObjectRef },
+    Raising { exception: PyBaseExceptionRef },
 
     // NoWorries,
     /// We are unwinding blocks, since we hit break
@@ -176,32 +177,16 @@ impl Frame {
                     // 1. Extract traceback from exception's '__traceback__' attr.
                     // 2. Add new entry with current execution position (filename, lineno, code_object) to traceback.
                     // 3. Unwind block stack till appropriate handler is found.
-                    assert!(objtype::isinstance(
-                        &exception,
-                        &vm.ctx.exceptions.base_exception_type
-                    ));
-
-                    let traceback = vm
-                        .get_attribute(exception.clone(), "__traceback__")
-                        .unwrap();
 
-                    let next = if vm.is_none(&traceback) {
-                        None
-                    } else {
-                        let traceback: PyTracebackRef = traceback
-                            .downcast()
-                            .expect("next must be a traceback object");
-                        Some(traceback)
-                    };
+                    let next = exception.traceback();
 
                     let new_traceback = PyTraceback::new(
                         next,
                         self.clone().into_ref(vm),
                         self.lasti.get(),
                         lineno.row(),
                     );
-                    vm.set_attr(&exception, "__traceback__", new_traceback.into_ref(vm))
-                        .unwrap();
+                    exception.set_traceback(Some(new_traceback.into_ref(vm)));
                     vm_trace!("Adding to traceback: {:?} {:?}", new_traceback, lineno);
 
                     match self.unwind_blocks(vm, UnwindReason::Raising { exception }) {
@@ -464,8 +449,8 @@ impl Frame {
                 let args = if let Some(exc) = exc {
                     let exc_type = exc.class().into_object();
                     let exc_val = exc.clone();
-                    let exc_tb = vm.ctx.none(); // TODO: retrieve traceback?
-                    vec![exc_type, exc_val, exc_tb]
+                    let exc_tb = exc.traceback().map_or(vm.get_none(), |tb| tb.into_object());
+                    vec![exc_type, exc_val.into_object(), exc_tb]
                 } else {
                     vec![vm.ctx.none(), vm.ctx.none(), vm.ctx.none()]
                 };
@@ -745,7 +730,7 @@ impl Frame {
                     self.pop_block();
                     if let UnwindReason::Raising { exception } = &reason {
                         self.push_block(BlockType::ExceptHandler {});
-                        self.push_value(exception.clone());
+                        self.push_value(exception.clone().into_object());
                         vm.push_exception(exception.clone());
                         self.jump(handler);
                         return Ok(None);
@@ -978,8 +963,8 @@ impl Frame {
 
     fn execute_raise(&self, vm: &VirtualMachine, argc: usize) -> FrameResult {
         let cause = match argc {
-            2 => self.get_exception(vm, true)?,
-            _ => vm.get_none(),
+            2 => Some(self.get_exception(vm, true)?),
+            _ => None,
         };
         let exception = match argc {
             0 => match vm.current_exception() {
@@ -996,18 +981,15 @@ impl Frame {
             _ => panic!("Invalid parameter for RAISE_VARARGS, must be between 0 to 3"),
         };
         let context = match argc {
-            0 => vm.get_none(), // We have already got the exception,
-            _ => match vm.current_exception() {
-                Some(exc) => exc,
-                None => vm.get_none(),
-            },
+            0 => None, // We have already got the exception,
+            _ => vm.current_exception(),
         };
         info!(
             "Exception raised: {:?} with cause: {:?} and context: {:?}",
             exception, cause, context
         );
-        vm.set_attr(&exception, vm.new_str("__cause__".to_string()), cause)?;
-        vm.set_attr(&exception, vm.new_str("__context__".to_string()), context)?;
+        exception.set_cause(cause);
+        exception.set_context(context);
         Err(exception)
     }
 
@@ -1385,12 +1367,16 @@ impl Frame {
     }
 
     #[cfg_attr(feature = "flame-it", flame("Frame"))]
-    fn get_exception(&self, vm: &VirtualMachine, none_allowed: bool) -> PyResult {
+    fn get_exception(
+        &self,
+        vm: &VirtualMachine,
+        none_allowed: bool,
+    ) -> PyResult<PyBaseExceptionRef> {
         let exception = self.pop_value();
-        if none_allowed && vm.get_none().is(&exception)
+        if none_allowed && vm.is_none(&exception)
             || objtype::isinstance(&exception, &vm.ctx.exceptions.base_exception_type)
         {
-            Ok(exception)
+            Ok(exception.downcast().unwrap())
         } else if let Ok(exc_type) = PyClassRef::try_from_object(vm, exception) {
             if objtype::issubclass(&exc_type, &vm.ctx.exceptions.base_exception_type) {
                 let exception = vm.new_empty_exception(exc_type)?;