@@ -1359,5 +1359,45 @@ def test__find_store_names(self):
13591359 self .assertEqual (res , expected )
13601360
13611361
1362+ class TestDisTraceback (unittest .TestCase ):
1363+ def setUp (self ) -> None :
1364+ try : # We need to clean up existing tracebacks
1365+ del sys .last_traceback
1366+ except AttributeError :
1367+ pass
1368+ return super ().setUp ()
1369+
1370+ def get_disassembly (self , tb ):
1371+ output = io .StringIO ()
1372+ with contextlib .redirect_stdout (output ):
1373+ dis .distb (tb )
1374+ return output .getvalue ()
1375+
1376+ def test_distb_empty (self ):
1377+ with self .assertRaises (RuntimeError ):
1378+ dis .distb ()
1379+
1380+ def test_distb_last_traceback (self ):
1381+ # We need to have an existing last traceback in `sys`:
1382+ tb = get_tb ()
1383+ sys .last_traceback = tb
1384+
1385+ self .assertEqual (self .get_disassembly (None ), dis_traceback )
1386+
1387+ def test_distb_explicit_arg (self ):
1388+ tb = get_tb ()
1389+
1390+ self .assertEqual (self .get_disassembly (tb ), dis_traceback )
1391+
1392+
1393+ class TestDisTracebackWithFile (TestDisTraceback ):
1394+ # Run the `distb` tests again, using the file arg instead of print
1395+ def get_disassembly (self , tb ):
1396+ output = io .StringIO ()
1397+ with contextlib .redirect_stdout (output ):
1398+ dis .distb (tb , file = output )
1399+ return output .getvalue ()
1400+
1401+
13621402if __name__ == "__main__" :
13631403 unittest .main ()
0 commit comments