1 files changed, 48 insertions, 4 deletions
diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
index 5bce2af68a..66a12bf12c 100644
--- a/Lib/test/test_traceback.py
+++ b/Lib/test/test_traceback.py
@@ -150,30 +150,74 @@ class SyntaxTracebackCases(unittest.TestCase):
 
 class TracebackFormatTests(unittest.TestCase):
 
-    def test_traceback_format(self):
+    def some_exception(self):
+        raise KeyError('blah')
+
+    def check_traceback_format(self, cleanup_func=None):
         try:
-            raise KeyError('blah')
+            self.some_exception()
         except KeyError:
             type_, value, tb = sys.exc_info()
+            if cleanup_func is not None:
+                # Clear the inner frames, not this one
+                cleanup_func(tb.tb_next)
             traceback_fmt = 'Traceback (most recent call last):\n' + \
                             ''.join(traceback.format_tb(tb))
             file_ = StringIO()
             traceback_print(tb, file_)
             python_fmt  = file_.getvalue()
+            # Call all _tb and _exc functions
+            with captured_output("stderr") as tbstderr:
+                traceback.print_tb(tb)
+            tbfile = StringIO()
+            traceback.print_tb(tb, file=tbfile)
+            with captured_output("stderr") as excstderr:
+                traceback.print_exc()
+            excfmt = traceback.format_exc()
+            excfile = StringIO()
+            traceback.print_exc(file=excfile)
         else:
             raise Error("unable to create test traceback string")
 
         # Make sure that Python and the traceback module format the same thing
         self.assertEqual(traceback_fmt, python_fmt)
+        # Now verify the _tb func output
+        self.assertEqual(tbstderr.getvalue(), tbfile.getvalue())
+        # Now verify the _exc func output
+        self.assertEqual(excstderr.getvalue(), excfile.getvalue())
+        self.assertEqual(excfmt, excfile.getvalue())
 
         # Make sure that the traceback is properly indented.
         tb_lines = python_fmt.splitlines()
-        self.assertEqual(len(tb_lines), 3)
-        banner, location, source_line = tb_lines
+        self.assertEqual(len(tb_lines), 5)
+        banner = tb_lines[0]
+        location, source_line = tb_lines[-2:]
         self.assertTrue(banner.startswith('Traceback'))
         self.assertTrue(location.startswith('  File'))
         self.assertTrue(source_line.startswith('    raise'))
 
+    def test_traceback_format(self):
+        self.check_traceback_format()
+
+    def test_traceback_format_with_cleared_frames(self):
+        # Check that traceback formatting also works with a clear()ed frame
+        def cleanup_tb(tb):
+            tb.tb_frame.clear()
+        self.check_traceback_format(cleanup_tb)
+
+    def test_stack_format(self):
+        # Verify _stack functions. Note we have to use _getframe(1) to
+        # compare them without this frame appearing in the output
+        with captured_output("stderr") as ststderr:
+            traceback.print_stack(sys._getframe(1))
+        stfile = StringIO()
+        traceback.print_stack(sys._getframe(1), file=stfile)
+        self.assertEqual(ststderr.getvalue(), stfile.getvalue())
+
+        stfmt = traceback.format_stack(sys._getframe(1))
+
+        self.assertEqual(ststderr.getvalue(), "".join(stfmt))
+
 
 cause_message = (
     "\nThe above exception was the direct cause "