From 74b6d7c0f5c5b5f2220cb09df1d96baa88e5c1f9 Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Sat, 2 Nov 2019 12:23:57 +0200
Subject: [PATCH] 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 @@ try:
 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<PyTracebackRef>,
     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)
 }