From f0c819acd0f85eafe12a7ff706650cb39d3fbf34 Mon Sep 17 00:00:00 2001
From: Robert Collins <rbtcollins@hp.com>
Date: Fri, 6 Mar 2015 13:46:35 +1300
Subject: Issue #22936: Allow showing local variables in unittest errors.

---
 Doc/library/unittest.rst          | 25 +++++++++++++++++------
 Lib/unittest/main.py              | 23 +++++++++++++++------
 Lib/unittest/result.py            |  7 +++++--
 Lib/unittest/runner.py            | 10 ++++++++-
 Lib/unittest/test/test_break.py   |  3 +++
 Lib/unittest/test/test_program.py | 26 +++++++++++++++++++----
 Lib/unittest/test/test_result.py  | 43 +++++++++++++++++++++++++++++++--------
 Lib/unittest/test/test_runner.py  | 10 ++++++---
 Misc/NEWS                         |  3 ++-
 9 files changed, 118 insertions(+), 32 deletions(-)

diff --git a/Doc/library/unittest.rst b/Doc/library/unittest.rst
index 92609ec..7ddf703 100644
--- a/Doc/library/unittest.rst
+++ b/Doc/library/unittest.rst
@@ -223,9 +223,16 @@ Command-line options
 
    Stop the test run on the first error or failure.
 
+.. cmdoption:: --locals
+
+   Show local variables in tracebacks.
+
 .. versionadded:: 3.2
    The command-line options ``-b``, ``-c`` and ``-f`` were added.
 
+.. versionadded:: 3.5
+   The command-line option ``--locals``.
+
 The command line can also be used for test discovery, for running all of the
 tests in a project or just a subset.
 
@@ -1782,12 +1789,10 @@ Loading and running tests
 
       Set to ``True`` when the execution of tests should stop by :meth:`stop`.
 
-
    .. attribute:: testsRun
 
       The total number of tests run so far.
 
-
    .. attribute:: buffer
 
       If set to true, ``sys.stdout`` and ``sys.stderr`` will be buffered in between
@@ -1797,7 +1802,6 @@ Loading and running tests
 
       .. versionadded:: 3.2
 
-
    .. attribute:: failfast
 
       If set to true :meth:`stop` will be called on the first failure or error,
@@ -1805,6 +1809,11 @@ Loading and running tests
 
       .. versionadded:: 3.2
 
+   .. attribute:: tb_locals
+
+      If set to true then local variables will be shown in tracebacks.
+
+      .. versionadded:: 3.5
 
    .. method:: wasSuccessful()
 
@@ -1815,7 +1824,6 @@ Loading and running tests
          Returns ``False`` if there were any :attr:`unexpectedSuccesses`
          from tests marked with the :func:`expectedFailure` decorator.
 
-
    .. method:: stop()
 
       This method can be called to signal that the set of tests being run should
@@ -1947,12 +1955,14 @@ Loading and running tests
 
 
 .. class:: TextTestRunner(stream=None, descriptions=True, verbosity=1, failfast=False, \
-                          buffer=False, resultclass=None, warnings=None)
+                          buffer=False, resultclass=None, warnings=None, *, tb_locals=False)
 
    A basic test runner implementation that outputs results to a stream. If *stream*
    is ``None``, the default, :data:`sys.stderr` is used as the output stream. This class
    has a few configurable parameters, but is essentially very simple.  Graphical
-   applications which run test suites should provide alternate implementations.
+   applications which run test suites should provide alternate implementations. Such
+   implementations should accept ``**kwargs`` as the interface to construct runners
+   changes when features are added to unittest.
 
    By default this runner shows :exc:`DeprecationWarning`,
    :exc:`PendingDeprecationWarning`, :exc:`ResourceWarning` and
@@ -1971,6 +1981,9 @@ Loading and running tests
       The default stream is set to :data:`sys.stderr` at instantiation time rather
       than import time.
 
+   .. versionchanged:: 3.5
+      Added the tb_locals parameter.
+
    .. method:: _makeResult()
 
       This method returns the instance of ``TestResult`` used by :meth:`run`.
diff --git a/Lib/unittest/main.py b/Lib/unittest/main.py
index 486d39f..b209a3a 100644
--- a/Lib/unittest/main.py
+++ b/Lib/unittest/main.py
@@ -58,7 +58,7 @@ class TestProgram(object):
     def __init__(self, module='__main__', defaultTest=None, argv=None,
                     testRunner=None, testLoader=loader.defaultTestLoader,
                     exit=True, verbosity=1, failfast=None, catchbreak=None,
-                    buffer=None, warnings=None):
+                    buffer=None, warnings=None, *, tb_locals=False):
         if isinstance(module, str):
             self.module = __import__(module)
             for part in module.split('.')[1:]:
@@ -73,6 +73,7 @@ class TestProgram(object):
         self.catchbreak = catchbreak
         self.verbosity = verbosity
         self.buffer = buffer
+        self.tb_locals = tb_locals
         if warnings is None and not sys.warnoptions:
             # even if DeprecationWarnings are ignored by default
             # print them anyway unless other warnings settings are
@@ -159,7 +160,9 @@ class TestProgram(object):
         parser.add_argument('-q', '--quiet', dest='verbosity',
                             action='store_const', const=0,
                             help='Quiet output')
-
+        parser.add_argument('--locals', dest='tb_locals',
+                            action='store_true',
+                            help='Show local variables in tracebacks')
         if self.failfast is None:
             parser.add_argument('-f', '--failfast', dest='failfast',
                                 action='store_true',
@@ -231,10 +234,18 @@ class TestProgram(object):
             self.testRunner = runner.TextTestRunner
         if isinstance(self.testRunner, type):
             try:
-                testRunner = self.testRunner(verbosity=self.verbosity,
-                                             failfast=self.failfast,
-                                             buffer=self.buffer,
-                                             warnings=self.warnings)
+                try:
+                    testRunner = self.testRunner(verbosity=self.verbosity,
+                                                 failfast=self.failfast,
+                                                 buffer=self.buffer,
+                                                 warnings=self.warnings,
+                                                 tb_locals=self.tb_locals)
+                except TypeError:
+                    # didn't accept the tb_locals argument
+                    testRunner = self.testRunner(verbosity=self.verbosity,
+                                                 failfast=self.failfast,
+                                                 buffer=self.buffer,
+                                                 warnings=self.warnings)
             except TypeError:
                 # didn't accept the verbosity, buffer or failfast arguments
                 testRunner = self.testRunner()
diff --git a/Lib/unittest/result.py b/Lib/unittest/result.py
index 8e0a643..a18f11b 100644
--- a/Lib/unittest/result.py
+++ b/Lib/unittest/result.py
@@ -45,6 +45,7 @@ class TestResult(object):
         self.unexpectedSuccesses = []
         self.shouldStop = False
         self.buffer = False
+        self.tb_locals = False
         self._stdout_buffer = None
         self._stderr_buffer = None
         self._original_stdout = sys.stdout
@@ -179,9 +180,11 @@ class TestResult(object):
         if exctype is test.failureException:
             # Skip assert*() traceback levels
             length = self._count_relevant_tb_levels(tb)
-            msgLines = traceback.format_exception(exctype, value, tb, length)
         else:
-            msgLines = traceback.format_exception(exctype, value, tb)
+            length = None
+        tb_e = traceback.TracebackException(
+            exctype, value, tb, limit=length, capture_locals=self.tb_locals)
+        msgLines = list(tb_e.format())
 
         if self.buffer:
             output = sys.stdout.getvalue()
diff --git a/Lib/unittest/runner.py b/Lib/unittest/runner.py
index 28b8865..2112262 100644
--- a/Lib/unittest/runner.py
+++ b/Lib/unittest/runner.py
@@ -126,7 +126,13 @@ class TextTestRunner(object):
     resultclass = TextTestResult
 
     def __init__(self, stream=None, descriptions=True, verbosity=1,
-                 failfast=False, buffer=False, resultclass=None, warnings=None):
+                 failfast=False, buffer=False, resultclass=None, warnings=None,
+                 *, tb_locals=False):
+        """Construct a TextTestRunner.
+
+        Subclasses should accept **kwargs to ensure compatibility as the
+        interface changes.
+        """
         if stream is None:
             stream = sys.stderr
         self.stream = _WritelnDecorator(stream)
@@ -134,6 +140,7 @@ class TextTestRunner(object):
         self.verbosity = verbosity
         self.failfast = failfast
         self.buffer = buffer
+        self.tb_locals = tb_locals
         self.warnings = warnings
         if resultclass is not None:
             self.resultclass = resultclass
@@ -147,6 +154,7 @@ class TextTestRunner(object):
         registerResult(result)
         result.failfast = self.failfast
         result.buffer = self.buffer
+        result.tb_locals = self.tb_locals
         with warnings.catch_warnings():
             if self.warnings:
                 # if self.warnings is set, use it to filter all the warnings
diff --git a/Lib/unittest/test/test_break.py b/Lib/unittest/test/test_break.py
index 0bf1a22..2c75019 100644
--- a/Lib/unittest/test/test_break.py
+++ b/Lib/unittest/test/test_break.py
@@ -211,6 +211,7 @@ class TestBreak(unittest.TestCase):
                 self.verbosity = verbosity
                 self.failfast = failfast
                 self.catchbreak = catchbreak
+                self.tb_locals = False
                 self.testRunner = FakeRunner
                 self.test = test
                 self.result = None
@@ -221,6 +222,7 @@ class TestBreak(unittest.TestCase):
         self.assertEqual(FakeRunner.initArgs, [((), {'buffer': None,
                                                      'verbosity': verbosity,
                                                      'failfast': failfast,
+                                                     'tb_locals': False,
                                                      'warnings': None})])
         self.assertEqual(FakeRunner.runArgs, [test])
         self.assertEqual(p.result, result)
@@ -235,6 +237,7 @@ class TestBreak(unittest.TestCase):
         self.assertEqual(FakeRunner.initArgs, [((), {'buffer': None,
                                                      'verbosity': verbosity,
                                                      'failfast': failfast,
+                                                     'tb_locals': False,
                                                      'warnings': None})])
         self.assertEqual(FakeRunner.runArgs, [test])
         self.assertEqual(p.result, result)
diff --git a/Lib/unittest/test/test_program.py b/Lib/unittest/test/test_program.py
index 725d67f..1cfc179 100644
--- a/Lib/unittest/test/test_program.py
+++ b/Lib/unittest/test/test_program.py
@@ -134,6 +134,7 @@ class InitialisableProgram(unittest.TestProgram):
     result = None
     verbosity = 1
     defaultTest = None
+    tb_locals = False
     testRunner = None
     testLoader = unittest.defaultTestLoader
     module = '__main__'
@@ -147,18 +148,19 @@ RESULT = object()
 class FakeRunner(object):
     initArgs = None
     test = None
-    raiseError = False
+    raiseError = 0
 
     def __init__(self, **kwargs):
         FakeRunner.initArgs = kwargs
         if FakeRunner.raiseError:
-            FakeRunner.raiseError = False
+            FakeRunner.raiseError -= 1
             raise TypeError
 
     def run(self, test):
         FakeRunner.test = test
         return RESULT
 
+
 class TestCommandLineArgs(unittest.TestCase):
 
     def setUp(self):
@@ -166,7 +168,7 @@ class TestCommandLineArgs(unittest.TestCase):
         self.program.createTests = lambda: None
         FakeRunner.initArgs = None
         FakeRunner.test = None
-        FakeRunner.raiseError = False
+        FakeRunner.raiseError = 0
 
     def testVerbosity(self):
         program = self.program
@@ -256,6 +258,7 @@ class TestCommandLineArgs(unittest.TestCase):
         self.assertEqual(FakeRunner.initArgs, {'verbosity': 'verbosity',
                                                 'failfast': 'failfast',
                                                 'buffer': 'buffer',
+                                                'tb_locals': False,
                                                 'warnings': 'warnings'})
         self.assertEqual(FakeRunner.test, 'test')
         self.assertIs(program.result, RESULT)
@@ -274,10 +277,25 @@ class TestCommandLineArgs(unittest.TestCase):
         self.assertEqual(FakeRunner.test, 'test')
         self.assertIs(program.result, RESULT)
 
+    def test_locals(self):
+        program = self.program
+
+        program.testRunner = FakeRunner
+        program.parseArgs([None, '--locals'])
+        self.assertEqual(True, program.tb_locals)
+        program.runTests()
+        self.assertEqual(FakeRunner.initArgs, {'buffer': False,
+                                               'failfast': False,
+                                               'tb_locals': True,
+                                               'verbosity': 1,
+                                               'warnings': None})
+
     def testRunTestsOldRunnerClass(self):
         program = self.program
 
-        FakeRunner.raiseError = True
+        # Two TypeErrors are needed to fall all the way back to old-style
+        # runners - one to fail tb_locals, one to fail buffer etc.
+        FakeRunner.raiseError = 2
         program.testRunner = FakeRunner
         program.verbosity = 'verbosity'
         program.failfast = 'failfast'
diff --git a/Lib/unittest/test/test_result.py b/Lib/unittest/test/test_result.py
index 489fe17..e39e2ea 100644
--- a/Lib/unittest/test/test_result.py
+++ b/Lib/unittest/test/test_result.py
@@ -8,6 +8,20 @@ import traceback
 import unittest
 
 
+class MockTraceback(object):
+    class TracebackException:
+        def __init__(self, *args, **kwargs):
+            self.capture_locals = kwargs.get('capture_locals', False)
+        def format(self):
+            result = ['A traceback']
+            if self.capture_locals:
+                result.append('locals')
+            return result
+
+def restore_traceback():
+    unittest.result.traceback = traceback
+
+
 class Test_TestResult(unittest.TestCase):
     # Note: there are not separate tests for TestResult.wasSuccessful(),
     # TestResult.errors, TestResult.failures, TestResult.testsRun or
@@ -227,6 +241,25 @@ class Test_TestResult(unittest.TestCase):
         self.assertIs(test_case, test)
         self.assertIsInstance(formatted_exc, str)
 
+    def test_addError_locals(self):
+        class Foo(unittest.TestCase):
+            def test_1(self):
+                1/0
+
+        test = Foo('test_1')
+        result = unittest.TestResult()
+        result.tb_locals = True
+
+        unittest.result.traceback = MockTraceback
+        self.addCleanup(restore_traceback)
+        result.startTestRun()
+        test.run(result)
+        result.stopTestRun()
+
+        self.assertEqual(len(result.errors), 1)
+        test_case, formatted_exc = result.errors[0]
+        self.assertEqual('A tracebacklocals', formatted_exc)
+
     def test_addSubTest(self):
         class Foo(unittest.TestCase):
             def test_1(self):
@@ -398,6 +431,7 @@ def __init__(self, stream=None, descriptions=None, verbosity=None):
     self.testsRun = 0
     self.shouldStop = False
     self.buffer = False
+    self.tb_locals = False
 
 classDict['__init__'] = __init__
 OldResult = type('OldResult', (object,), classDict)
@@ -454,15 +488,6 @@ class Test_OldTestResult(unittest.TestCase):
         runner.run(Test('testFoo'))
 
 
-class MockTraceback(object):
-    @staticmethod
-    def format_exception(*_):
-        return ['A traceback']
-
-def restore_traceback():
-    unittest.result.traceback = traceback
-
-
 class TestOutputBuffering(unittest.TestCase):
 
     def setUp(self):
diff --git a/Lib/unittest/test/test_runner.py b/Lib/unittest/test/test_runner.py
index 7c0bd51..9cbc260 100644
--- a/Lib/unittest/test/test_runner.py
+++ b/Lib/unittest/test/test_runner.py
@@ -158,7 +158,7 @@ class Test_TextTestRunner(unittest.TestCase):
         self.assertEqual(runner.warnings, None)
         self.assertTrue(runner.descriptions)
         self.assertEqual(runner.resultclass, unittest.TextTestResult)
-
+        self.assertFalse(runner.tb_locals)
 
     def test_multiple_inheritance(self):
         class AResult(unittest.TestResult):
@@ -172,14 +172,13 @@ class Test_TextTestRunner(unittest.TestCase):
         # on arguments in its __init__ super call
         ATextResult(None, None, 1)
 
-
     def testBufferAndFailfast(self):
         class Test(unittest.TestCase):
             def testFoo(self):
                 pass
         result = unittest.TestResult()
         runner = unittest.TextTestRunner(stream=io.StringIO(), failfast=True,
-                                           buffer=True)
+                                         buffer=True)
         # Use our result object
         runner._makeResult = lambda: result
         runner.run(Test('testFoo'))
@@ -187,6 +186,11 @@ class Test_TextTestRunner(unittest.TestCase):
         self.assertTrue(result.failfast)
         self.assertTrue(result.buffer)
 
+    def test_locals(self):
+        runner = unittest.TextTestRunner(stream=io.StringIO(), tb_locals=True)
+        result = runner.run(unittest.TestSuite())
+        self.assertEqual(True, result.tb_locals)
+
     def testRunnerRegistersResult(self):
         class Test(unittest.TestCase):
             def testFoo(self):
diff --git a/Misc/NEWS b/Misc/NEWS
index 9982079..1294a4d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -39,7 +39,8 @@ Library
 - Issue #21619: Popen objects no longer leave a zombie after exit in the with
   statement if the pipe was broken.  Patch by Martin Panter.
 
-- Issue #22936: Make it possible to show local variables in tracebacks.
+- Issue #22936: Make it possible to show local variables in tracebacks for
+  both the traceback module and unittest.
 
 - Issue #15955: Add an option to limit the output size in bz2.decompress().
   Patch by Nikolaus Rath.
-- 
cgit v0.12

