@@ -1359,5 +1359,45 @@ def test__find_store_names(self):
1359
1359
self .assertEqual (res , expected )
1360
1360
1361
1361
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
+
1362
1402
if __name__ == "__main__" :
1363
1403
unittest .main ()
0 commit comments