From e8ed336c28bd96071730161e2b5262b6c3b2c4b5 Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 26 Oct 2019 12:42:58 +0300 Subject: [PATCH 1/3] Add with_traceback for exception base --- vm/src/exceptions.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs index 79924df9c8..38ec24b4ca 100644 --- a/vm/src/exceptions.rs +++ b/vm/src/exceptions.rs @@ -208,6 +208,15 @@ fn exception_repr(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { Ok(vm.new_str(joined_str)) } +fn exception_with_traceback( + zelf: PyObjectRef, + tb: PyTracebackRef, + vm: &VirtualMachine, +) -> PyResult { + vm.set_attr(&zelf, "__traceback__", tb)?; + Ok(zelf) +} + #[derive(Debug)] pub struct ExceptionZoo { pub arithmetic_error: PyClassRef, @@ -400,7 +409,8 @@ fn import_error_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { pub fn init(context: &PyContext) { let base_exception_type = &context.exceptions.base_exception_type; extend_class!(context, base_exception_type, { - "__init__" => context.new_rustfunc(exception_init) + "__init__" => context.new_rustfunc(exception_init), + "with_traceback" => context.new_rustfunc(exception_with_traceback) }); let exception_type = &context.exceptions.exception_type; From 1fb868a08aa210db688b25fcf4b09ba8c9c9aa15 Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Fri, 1 Nov 2019 14:29:58 +0200 Subject: [PATCH 2/3] Add test for exception with_traceback --- tests/snippets/stdlib_traceback.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/snippets/stdlib_traceback.py b/tests/snippets/stdlib_traceback.py index 12ac9835ec..808011fc6b 100644 --- a/tests/snippets/stdlib_traceback.py +++ b/tests/snippets/stdlib_traceback.py @@ -3,5 +3,15 @@ try: 1/0 except ZeroDivisionError as ex: - tb = traceback.format_list(traceback.extract_tb(ex.__traceback__)) + tb = traceback.extract_tb(ex.__traceback__) assert len(tb) == 1 + + +try: + try: + 1/0 + except ZeroDivisionError as ex: + raise KeyError().with_traceback(ex.__traceback__) +except KeyError as ex2: + tb = traceback.extract_tb(ex2.__traceback__) + assert tb[1].line == "1/0" From 74b6d7c0f5c5b5f2220cb09df1d96baa88e5c1f9 Mon Sep 17 00:00:00 2001 From: Aviv Palivoda Date: Sat, 2 Nov 2019 12:23:57 +0200 Subject: [PATCH 3/3] Allow None as with_tracebak arg --- tests/snippets/stdlib_traceback.py | 10 ++++++++++ vm/src/exceptions.rs | 8 ++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/snippets/stdlib_traceback.py b/tests/snippets/stdlib_traceback.py index 808011fc6b..689f36e027 100644 --- a/tests/snippets/stdlib_traceback.py +++ b/tests/snippets/stdlib_traceback.py @@ -15,3 +15,13 @@ except KeyError as ex2: tb = traceback.extract_tb(ex2.__traceback__) assert tb[1].line == "1/0" + + +try: + try: + 1/0 + except ZeroDivisionError as ex: + raise ex.with_traceback(None) +except ZeroDivisionError as ex2: + tb = traceback.extract_tb(ex2.__traceback__) + assert len(tb) == 1 diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs index 38ec24b4ca..43682a8d9d 100644 --- a/vm/src/exceptions.rs +++ b/vm/src/exceptions.rs @@ -210,10 +210,14 @@ fn exception_repr(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { fn exception_with_traceback( zelf: PyObjectRef, - tb: PyTracebackRef, + tb: Option, vm: &VirtualMachine, ) -> PyResult { - vm.set_attr(&zelf, "__traceback__", tb)?; + vm.set_attr( + &zelf, + "__traceback__", + tb.map_or(vm.get_none(), |tb| tb.into_object()), + )?; Ok(zelf) }