Merge pull request #2619 from blueyed/pdb-resume-capture

pdb: resume capturing after `continue`

changelog/2619.feature.rst

@@ -0,0 +1,4 @@
+Resume capturing output after ``continue`` with ``__import__("pdb").set_trace()``.
+
+This also adds a new ``pytest_leave_pdb`` hook, and passes in ``pdb`` to the
+existing ``pytest_enter_pdb`` hook.

src/_pytest/debugging.py

@@ -95,9 +95,47 @@ def set_trace(cls, set_break=True):
             tw = _pytest.config.create_terminal_writer(cls._config)
             tw.line()
             tw.sep(">", "PDB set_trace (IO-capturing turned off)")
-            cls._pluginmanager.hook.pytest_enter_pdb(config=cls._config)
+
+            class _PdbWrapper(cls._pdb_cls, object):
+                _pytest_capman = capman
+                _continued = False
+
+                def do_continue(self, arg):
+                    ret = super(_PdbWrapper, self).do_continue(arg)
+                    if self._pytest_capman:
+                        tw = _pytest.config.create_terminal_writer(cls._config)
+                        tw.line()
+                        tw.sep(">", "PDB continue (IO-capturing resumed)")
+                        self._pytest_capman.resume_global_capture()
+                    cls._pluginmanager.hook.pytest_leave_pdb(
+                        config=cls._config, pdb=self
+                    )
+                    self._continued = True
+                    return ret
+
+                do_c = do_cont = do_continue
+
+                def setup(self, f, tb):
+                    """Suspend on setup().
+
+                    Needed after do_continue resumed, and entering another
+                    breakpoint again.
+                    """
+                    ret = super(_PdbWrapper, self).setup(f, tb)
+                    if not ret and self._continued:
+                        # pdb.setup() returns True if the command wants to exit
+                        # from the interaction: do not suspend capturing then.
+                        if self._pytest_capman:
+                            self._pytest_capman.suspend_global_capture(in_=True)
+                    return ret
+
+            _pdb = _PdbWrapper()
+            cls._pluginmanager.hook.pytest_enter_pdb(config=cls._config, pdb=_pdb)
+        else:
+            _pdb = cls._pdb_cls()
+
         if set_break:
-            cls._pdb_cls().set_trace(frame)
+            _pdb.set_trace(frame)
 
 
 class PdbInvoke(object):

src/_pytest/hookspec.py

@@ -603,9 +603,21 @@ def pytest_exception_interact(node, call, report):
     """
 
 
-def pytest_enter_pdb(config):
+def pytest_enter_pdb(config, pdb):
     """ called upon pdb.set_trace(), can be used by plugins to take special
     action just before the python debugger enters in interactive mode.
 
     :param _pytest.config.Config config: pytest config object
+    :param pdb.Pdb pdb: Pdb instance
+    """
+
+
+def pytest_leave_pdb(config, pdb):
+    """ called when leaving pdb (e.g. with continue after pdb.set_trace()).
+
+    Can be used by plugins to take special action just after the python
+    debugger leaves interactive mode.
+
+    :param _pytest.config.Config config: pytest config object
+    :param pdb.Pdb pdb: Pdb instance
     """

testing/test_pdb.py

@@ -158,6 +158,7 @@ def test_not_called_due_to_quit():
         assert "= 1 failed in" in rest
         assert "def test_1" not in rest
         assert "Exit: Quitting debugger" in rest
+        assert "PDB continue (IO-capturing resumed)" not in rest
         self.flush(child)
 
     @staticmethod
@@ -489,18 +490,23 @@ def test_1():
         """
         )
         child = testdir.spawn_pytest(str(p1))
+        child.expect(r"PDB set_trace \(IO-capturing turned off\)")
         child.expect("test_1")
         child.expect("x = 3")
         child.expect("Pdb")
         child.sendline("c")
+        child.expect(r"PDB continue \(IO-capturing resumed\)")
+        child.expect(r"PDB set_trace \(IO-capturing turned off\)")
         child.expect("x = 4")
         child.expect("Pdb")
         child.sendeof()
+        child.expect("_ test_1 _")
+        child.expect("def test_1")
+        child.expect("Captured stdout call")
         rest = child.read().decode("utf8")
-        assert "1 failed" in rest
-        assert "def test_1" in rest
         assert "hello17" in rest  # out is captured
         assert "hello18" in rest  # out is captured
+        assert "1 failed" in rest
         self.flush(child)
 
     def test_pdb_used_outside_test(self, testdir):
@@ -541,15 +547,29 @@ def test_pdb_collection_failure_is_shown(self, testdir):
             ["E   NameError: *xxx*", "*! *Exit: Quitting debugger !*"]  # due to EOF
         )
 
-    def test_enter_pdb_hook_is_called(self, testdir):
+    def test_enter_leave_pdb_hooks_are_called(self, testdir):
         testdir.makeconftest(
             """
-            def pytest_enter_pdb(config):
-                assert config.testing_verification == 'configured'
-                print('enter_pdb_hook')
+            mypdb = None
 
             def pytest_configure(config):
                 config.testing_verification = 'configured'
+
+            def pytest_enter_pdb(config, pdb):
+                assert config.testing_verification == 'configured'
+                print('enter_pdb_hook')
+
+                global mypdb
+                mypdb = pdb
+                mypdb.set_attribute = "bar"
+
+            def pytest_leave_pdb(config, pdb):
+                assert config.testing_verification == 'configured'
+                print('leave_pdb_hook')
+
+                global mypdb
+                assert mypdb is pdb
+                assert mypdb.set_attribute == "bar"
         """
         )
         p1 = testdir.makepyfile(
@@ -558,11 +578,17 @@ def pytest_configure(config):
 
             def test_foo():
                 pytest.set_trace()
+                assert 0
         """
         )
         child = testdir.spawn_pytest(str(p1))
         child.expect("enter_pdb_hook")
-        child.send("c\n")
+        child.sendline("c")
+        child.expect(r"PDB continue \(IO-capturing resumed\)")
+        child.expect("Captured stdout call")
+        rest = child.read().decode("utf8")
+        assert "leave_pdb_hook" in rest
+        assert "1 failed" in rest
         child.sendeof()
         self.flush(child)