RustPython · youknowone · Aug 9, 2021 · Aug 9, 2021 · Aug 9, 2021
diff --git a/extra_tests/snippets/builtin_exceptions.py b/extra_tests/snippets/builtin_exceptions.py
@@ -0,0 +1,154 @@
+import platform
+import sys
+
+
+# Regression to:
+# https://github.com/RustPython/RustPython/issues/2779
+
+class MyError(Exception):
+    pass
+
+
+e = MyError('message')
+
+try:
+    raise e from e
+except MyError as exc:
+    # It was a segmentation fault before, will print info to stdout:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert isinstance(exc, MyError)
+    assert exc.__cause__ is e
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+try:
+    raise ValueError('test') from e
+except ValueError as exc:
+    sys.excepthook(type(exc), exc, exc.__traceback__)  # ok, will print two excs
+    assert isinstance(exc, ValueError)
+    assert exc.__cause__ is e
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+
+# New case:
+# potential recursion on `__context__` field
+
+e = MyError('message')
+
+try:
+    try:
+        raise e
+    except MyError as exc:
+        raise e
+    else:
+        assert False, 'exception not raised'
+except MyError as exc:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert exc.__cause__ is None
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+e = MyError('message')
+
+try:
+    try:
+        raise e
+    except MyError as exc:
+        raise exc
+    else:
+        assert False, 'exception not raised'
+except MyError as exc:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert exc.__cause__ is None
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+e = MyError('message')
+
+try:
+    try:
+        raise e
+    except MyError as exc:
+        raise e from e
+    else:
+        assert False, 'exception not raised'
+except MyError as exc:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert exc.__cause__ is e
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+e = MyError('message')
+
+try:
+    try:
+        raise e
+    except MyError as exc:
+        raise exc from e
+    else:
+        assert False, 'exception not raised'
+except MyError as exc:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert exc.__cause__ is e
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+
+# New case:
+# two exception in a recursion loop
+
+class SubError(MyError):
+    pass
+
+e = MyError('message')
+d = SubError('sub')
+
+
+try:
+    raise e from d
+except MyError as exc:
+    # It was a segmentation fault before, will print info to stdout:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert isinstance(exc, MyError)
+    assert exc.__cause__ is d
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+e = MyError('message')
+
+try:
+    raise d from e
+except SubError as exc:
+    # It was a segmentation fault before, will print info to stdout:
+    sys.excepthook(type(exc), exc, exc.__traceback__)
+    assert isinstance(exc, SubError)
+    assert exc.__cause__ is e
+    assert exc.__context__ is None
+else:
+    assert False, 'exception not raised'
+
+
+# New case:
+# explicit `__context__` manipulation.
+
+e = MyError('message')
+e.__context__ = e
+
+try:
+    raise e
+except MyError as exc:
+    # It was a segmentation fault before, will print info to stdout:
+    if platform.python_implementation() == 'RustPython':
+        # For some reason `CPython` hangs on this code:
+        sys.excepthook(type(exc), exc, exc.__traceback__)
+        assert isinstance(exc, MyError)
+        assert exc.__cause__ is None
+        assert exc.__context__ is e
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
@@ -11,12 +11,13 @@ use crate::types::create_type_with_slots;
 use crate::StaticType;
 use crate::VirtualMachine;
 use crate::{
-    IntoPyObject, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
-    TryFromObject, TypeProtocol,
+    IdProtocol, IntoPyObject, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult,
+    PyValue, TryFromObject, TypeProtocol,
 };
 
 use crossbeam_utils::atomic::AtomicCell;
 use itertools::Itertools;
+use std::collections::HashSet;
 use std::fmt;
 use std::fs::File;
 use std::io::{self, BufRead, BufReader};
@@ -194,21 +195,8 @@ pub fn write_exception<W: Write>(
     vm: &VirtualMachine,
     exc: &PyBaseExceptionRef,
 ) -> Result<(), W::Error> {
-    if let Some(cause) = exc.cause() {
-        write_exception(output, vm, &cause)?;
-        writeln!(
-            output,
-            "\nThe above exception was the direct cause of the following exception:\n"
-        )?;
-    } else if let Some(context) = exc.context() {
-        write_exception(output, vm, &context)?;
-        writeln!(
-            output,
-            "\nDuring handling of the above exception, another exception occurred:\n"
-        )?;
-    }
-
-    write_exception_inner(output, vm, exc)
+    let seen = &mut HashSet::<usize>::new();
+    write_exception_recursive(output, vm, exc, seen)
 }
 
 fn print_source_line<W: Write>(
@@ -253,6 +241,49 @@ fn write_traceback_entry<W: Write>(
     Ok(())
 }
 
+fn write_exception_recursive<W: Write>(
+    output: &mut W,
+    vm: &VirtualMachine,
+    exc: &PyBaseExceptionRef,
+    seen: &mut HashSet<usize>,
+) -> Result<(), W::Error> {
+    // This function should not be called directly,
+    // use `wite_exception` as a public interface.
+    // It is similar to `print_exception_recursive` from `CPython`.
+    seen.insert(exc.as_object().get_id());
+
+    #[allow(clippy::manual_map)]
+    if let Some((cause_or_context, msg)) = if let Some(cause) = exc.cause() {
+        // This can be a special case: `raise e from e`,
+        // we just ignore it and treat like `raise e` without any extra steps.
+        Some((
+            cause,
+            "\nThe above exception was the direct cause of the following exception:\n",
+        ))
+    } else if let Some(context) = exc.context() {
+        // This can be a special case:
+        //   e = ValueError('e')
+        //   e.__context__ = e
+        // In this case, we just ignore
+        // `__context__` part from going into recursion.
+        Some((
+            context,
+            "\nDuring handling of the above exception, another exception occurred:\n",
+        ))
+    } else {
+        None
+    } {
+        if !seen.contains(&cause_or_context.as_object().get_id()) {
+            write_exception_recursive(output, vm, &cause_or_context, seen)?;
+            writeln!(output, "{}", msg)?;
+        } else {
+            seen.insert(cause_or_context.as_object().get_id());
+        }
+    }
+
+    write_exception_inner(output, vm, exc)
+}
+
 /// Print exception with traceback
 pub fn write_exception_inner<W: Write>(
     output: &mut W,