Merge pull request RustPython#821 from palaviv/excpetion-cause

coolreader18 · web-flow · commit 22203975fadd · 2019-04-12T16:08:54.000-05:00
Add exception cause
diff --git a/tests/snippets/try_exceptions.py b/tests/snippets/try_exceptions.py
@@ -1,4 +1,4 @@
-
+from testutils import assertRaises
 
 try:
     raise BaseException()
@@ -95,3 +95,33 @@ def __init__(self):
     l.append(3)
     print('boom', type(ex))
 assert l == [1, 3]
+
+cause = None
+try:
+    try:
+        raise ZeroDivisionError
+    except ZeroDivisionError as ex:
+        assert ex.__cause__ == None
+        cause = ex
+        raise NameError from ex
+except NameError as ex2:
+    assert ex2.__cause__ == cause
+
+try:
+    raise ZeroDivisionError from None
+except ZeroDivisionError as ex:
+    assert ex.__cause__ == None
+
+try:
+    raise ZeroDivisionError
+except ZeroDivisionError as ex:
+    assert ex.__cause__ == None
+
+with assertRaises(TypeError):
+    raise ZeroDivisionError from 5
+
+try:
+    raise ZeroDivisionError from NameError
+except ZeroDivisionError as ex:
+    assert type(ex.__cause__) == NameError
+
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
@@ -2,7 +2,7 @@ use crate::function::PyFuncArgs;
 use crate::obj::objsequence;
 use crate::obj::objtype;
 use crate::obj::objtype::PyClassRef;
-use crate::pyobject::{create_type, PyContext, PyObjectRef, PyResult, TypeProtocol};
+use crate::pyobject::{create_type, IdProtocol, PyContext, PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 
 fn exception_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -18,8 +18,19 @@ fn exception_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
 
-// Print exception including traceback:
+// print excption chain
 pub fn print_exception(vm: &VirtualMachine, exc: &PyObjectRef) {
+    if let Ok(cause) = vm.get_attribute(exc.clone(), "__cause__") {
+        if !vm.get_none().is(&cause) {
+            print_exception(vm, &cause);
+            println!("\nThe above exception was the direct cause of the following exception:\n");
+        }
+    }
+    print_exception_inner(vm, exc)
+}
+
+// Print exception including traceback:
+pub fn print_exception_inner(vm: &VirtualMachine, exc: &PyObjectRef) {
     if let Ok(tb) = vm.get_attribute(exc.clone(), "__traceback__") {
         println!("Traceback (most recent call last):");
         if objtype::isinstance(&tb, &vm.ctx.list_type()) {
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -688,31 +688,18 @@ impl Frame {
             }
 
             bytecode::Instruction::Raise { argc } => {
+                let cause = match argc {
+                    2 => self.get_exception(vm, true)?,
+                    _ => vm.get_none(),
+                };
                 let exception = match argc {
-                    1 => self.pop_value(),
-                    0 | 2 | 3 => panic!("Not implemented!"),
+                    1 | 2 => self.get_exception(vm, false)?,
+                    0 | 3 => panic!("Not implemented!"),
                     _ => panic!("Invalid parameter for RAISE_VARARGS, must be between 0 to 3"),
                 };
-                if objtype::isinstance(&exception, &vm.ctx.exceptions.base_exception_type) {
-                    info!("Exception raised: {:?}", exception);
-                    Err(exception)
-                } else if let Ok(exception) = PyClassRef::try_from_object(vm, exception) {
-                    if objtype::issubclass(&exception, &vm.ctx.exceptions.base_exception_type) {
-                        let exception = vm.new_empty_exception(exception)?;
-                        info!("Exception raised: {:?}", exception);
-                        Err(exception)
-                    } else {
-                        let msg = format!(
-                            "Can only raise BaseException derived types, not {}",
-                            exception
-                        );
-                        let type_error_type = vm.ctx.exceptions.type_error.clone();
-                        let type_error = vm.new_exception(type_error_type, msg);
-                        Err(type_error)
-                    }
-                } else {
-                    Err(vm.new_type_error("exceptions must derive from BaseException".to_string()))
-                }
+                info!("Exception raised: {:?} with cause: {:?}", exception, cause);
+                vm.set_attr(&exception, vm.new_str("__cause__".to_string()), cause)?;
+                Err(exception)
             }
 
             bytecode::Instruction::Break => {
@@ -1215,6 +1202,28 @@ impl Frame {
         let stack = self.stack.borrow_mut();
         stack[stack.len() - depth - 1].clone()
     }
+
+    fn get_exception(&self, vm: &VirtualMachine, none_allowed: bool) -> PyResult {
+        let exception = self.pop_value();
+        if none_allowed && vm.get_none().is(&exception) {
+            Ok(exception)
+        } else if objtype::isinstance(&exception, &vm.ctx.exceptions.base_exception_type) {
+            Ok(exception)
+        } else if let Ok(exception) = PyClassRef::try_from_object(vm, exception) {
+            if objtype::issubclass(&exception, &vm.ctx.exceptions.base_exception_type) {
+                let exception = vm.new_empty_exception(exception)?;
+                Ok(exception)
+            } else {
+                let msg = format!(
+                    "Can only raise BaseException derived types, not {}",
+                    exception
+                );
+                Err(vm.new_type_error(msg))
+            }
+        } else {
+            Err(vm.new_type_error("exceptions must derive from BaseException".to_string()))
+        }
+    }
 }
 
 impl fmt::Debug for Frame {
