Have _PdbClient work with the socket directly

Previously we worked with a file object created with `socket.makefile`, but on Windows the `readline()` method of a socket file can't be interrupted with Ctrl-C, and neither can `recv()` on a socket. This refactor lays the ground work for a solution this this problem.
python · pablogsal · May 5, 2025 · Apr 30, 2025 · Apr 30, 2025 · Apr 30, 2025
commit f65f99ffb48cc472d193c21db019d7bf6b96edd6
@@ -2903,9 +2903,10 @@ def default(self, line):
 
 
 class _PdbClient:
-    def __init__(self, pid, sockfile, interrupt_script):
+    def __init__(self, pid, server_socket, interrupt_script):
         self.pid = pid
-        self.sockfile = sockfile
+        self.read_buf = b""
+        self.server_socket = server_socket
         self.interrupt_script = interrupt_script
         self.pdb_instance = Pdb()
         self.pdb_commands = set()
@@ -2947,8 +2948,7 @@ def _send(self, **kwargs):
         self._ensure_valid_message(kwargs)
         json_payload = json.dumps(kwargs)
         try:
-            self.sockfile.write(json_payload.encode() + b"\n")
-            self.sockfile.flush()
+            self.server_socket.sendall(json_payload.encode() + b"\n")
         except OSError:
             # This means that the client has abruptly disconnected, but we'll
             # handle that the next time we try to read from the client instead
@@ -2957,6 +2957,15 @@ def _send(self, **kwargs):
             # return an empty string because the socket may be half-closed.
             self.write_failed = True
 
+    def _readline(self):
+        while b"\n" not in self.read_buf:
+            self.read_buf += self.server_socket.recv(16 * 1024)
+            if not self.read_buf:
+                return b""
+
+        ret, sep, self.read_buf = self.read_buf.partition(b"\n")
+        return ret + sep
+
     def read_command(self, prompt):
         reply = input(prompt)
 
@@ -3054,7 +3063,7 @@ def cmdloop(self):
             while not self.write_failed:
                 try:
                     with self._handle_sigint(signal.default_int_handler):
-                        if not (payload_bytes := self.sockfile.readline()):
+                        if not (payload_bytes := self._readline()):
                             break
                 except KeyboardInterrupt:
                     self.send_interrupt()
@@ -3144,7 +3153,7 @@ def complete(self, text, state):
                 return None
 
             with self._handle_sigint(signal.default_int_handler):
-                payload = self.sockfile.readline()
+                payload = self._readline()
             if not payload:
                 return None
 
@@ -3211,9 +3220,6 @@ def attach(pid, commands=()):
             client_sock, _ = server.accept()
 
             with closing(client_sock):
-                sockfile = client_sock.makefile("rwb")
-
-            with closing(sockfile):
                 with tempfile.NamedTemporaryFile("w", delete_on_close=False) as interrupt_script:
                     interrupt_script.write(
                         'import pdb, sys\n'
@@ -3222,7 +3228,7 @@ def attach(pid, commands=()):
                     )
                     interrupt_script.close()
 
-                    _PdbClient(pid, sockfile, interrupt_script.name).cmdloop()
+                    _PdbClient(pid, client_sock, interrupt_script.name).cmdloop()
 
 
 # Post-Mortem interface

@@ -12,7 +12,7 @@
 import threading
 import unittest
 import unittest.mock
-from contextlib import contextmanager, redirect_stdout, ExitStack
+from contextlib import closing, contextmanager, redirect_stdout, ExitStack
 from pathlib import Path
 from test.support import is_wasi, os_helper, requires_subprocess, SHORT_TIMEOUT
 from test.support.os_helper import temp_dir, TESTFN, unlink
@@ -79,52 +79,14 @@ def get_output(self) -> List[dict]:
         return results
 
 
-class MockDebuggerSocket:
-    """Mock file-like simulating a connection to a _RemotePdb instance"""
-
-    def __init__(self, incoming):
-        self.incoming = iter(incoming)
-        self.outgoing = []
-        self.buffered = bytearray()
-
-    def write(self, data: bytes) -> None:
-        """Simulate write to socket."""
-        self.buffered += data
-
-    def flush(self) -> None:
-        """Ensure each line is valid JSON."""
-        lines = self.buffered.splitlines(keepends=True)
-        self.buffered.clear()
-        for line in lines:
-            assert line.endswith(b"\n")
-            self.outgoing.append(json.loads(line))
-
-    def readline(self) -> bytes:
-        """Read a line from the prepared input queue."""
-        # Anything written must be flushed before trying to read,
-        # since the read will be dependent upon the last write.
-        assert not self.buffered
-        try:
-            item = next(self.incoming)
-            if not isinstance(item, bytes):
-                item = json.dumps(item).encode()
-            return item + b"\n"
-        except StopIteration:
-            return b""
-
-    def close(self) -> None:
-        """No-op close implementation."""
-        pass
-
-
 class PdbClientTestCase(unittest.TestCase):
     """Tests for the _PdbClient class."""
 
     def do_test(
         self,
         *,
         incoming,
-        simulate_failure=None,
+        simulate_send_failure=False,
         expected_outgoing=None,
         expected_completions=None,
         expected_exception=None,
@@ -142,16 +104,6 @@ def do_test(
         expected_state.setdefault("write_failed", False)
         messages = [m for source, m in incoming if source == "server"]
         prompts = [m["prompt"] for source, m in incoming if source == "user"]
-        sockfile = MockDebuggerSocket(messages)
-        stdout = io.StringIO()
-
-        if simulate_failure:
-            sockfile.write = unittest.mock.Mock()
-            sockfile.flush = unittest.mock.Mock()
-            if simulate_failure == "write":
-                sockfile.write.side_effect = OSError(&quo
E864
t;write failed")
-            elif simulate_failure == "flush":
-                sockfile.flush.side_effect = OSError("flush failed")
 
         input_iter = (m for source, m in incoming if source == "user")
         completions = []
@@ -181,14 +133,36 @@ def mock_input(prompt):
             return reply
 
         with ExitStack() as stack:
+            client_sock, server_sock = socket.socketpair()
+            stack.enter_context(closing(client_sock))
+            stack.enter_context(closing(server_sock))
+
+            server_sock = unittest.mock.Mock(wraps=server_sock)
+
+            client_sock.sendall(
+                b"".join(
+                    (m if isinstance(m, bytes) else json.dumps(m).encode()) + b"\n"
+                    for m in messages
+                )
+            )
+            client_sock.shutdown(socket.SHUT_WR)
+
+            if simulate_send_failure:
+                server_sock.sendall = unittest.mock.Mock(
+                    side_effect=OSError("sendall failed")
+                )
+                client_sock.shutdown(socket.SHUT_RD)
+
+            stdout = io.StringIO()
+
             input_mock = stack.enter_context(
                 unittest.mock.patch("pdb.input", side_effect=mock_input)
             )
             stack.enter_context(redirect_stdout(stdout))
 
             client = _PdbClient(
                 pid=0,
-                sockfile=sockfile,
+                server_socket=server_sock,
                 interrupt_script="/a/b.py",
             )
 
@@ -199,13 +173,12 @@ def mock_input(prompt):
 
             client.cmdloop()
 
-        actual_outgoing = sockfile.outgoing
-        if simulate_failure:
-            actual_outgoing += [
-                json.loads(msg.args[0]) for msg in sockfile.write.mock_calls
-            ]
+        sent_msgs = [msg.args[0] for msg in server_sock.sendall.mock_calls]
+        for msg in sent_msgs:
+            assert msg.endswith(b"\n")
+        actual_outgoing = [json.loads(msg) for msg in sent_msgs]
 
-        self.assertEqual(sockfile.outgoing, expected_outgoing)
+        self.assertEqual(actual_outgoing, expected_outgoing)
         self.assertEqual(completions, expected_completions)
         if expected_stdout_substring and not expected_stdout:
             self.assertIn(expected_stdout_substring, stdout.getvalue())
@@ -478,20 +451,7 @@ def test_write_failing(self):
         self.do_test(
             incoming=incoming,
             expected_outgoing=[{"signal": "INT"}],
-            simulate_failure="write",
-            expected_state={"write_failed": True},
-        )
-
-    def test_flush_failing(self):
-        """Test terminating if flush fails due to a half closed socket."""
-        incoming = [
-            ("server", {"prompt": "(Pdb) ", "state": "pdb"}),
-            ("user", {"prompt": "(Pdb) ", "input": KeyboardInterrupt()}),
-        ]
-        self.do_test(
-            incoming=incoming,
-            expected_outgoing=[{"signal": "INT"}],
-            simulate_failure="flush",
+            simulate_send_failure=True,
             expected_state={"write_failed": True},
         )
 
@@ -622,42 +582,7 @@ def test_write_failure_during_completion(self):
                 },
                 {"reply": "xyz"},
             ],
-            simulate_failure="write",
-            expected_completions=[],
-            expected_state={"state": "interact", "write_failed": True},
-        )
-
-    def test_flush_failure_during_completion(self):
-        """Test failing to flush to the socket to request tab completions."""
-        incoming = [
-            ("server", {"prompt": ">>> ", "state": "interact"}),
-            (
-                "user",
-                {
-                    "prompt": ">>> ",
-                    "completion_request": {
-                        "line": "xy",
-                        "begidx": 0,
-                        "endidx": 2,
-                    },
-                    "input": "xyz",
-                },
-            ),
-        ]
-        self.do_test(
-            incoming=incoming,
-            expected_outgoing=[
-                {
-                    "complete": {
-                        "text": "xy",
-                        "line": "xy",
-                        "begidx": 0,
-                        "endidx": 2,
-                    }
-                },
-                {"reply": "xyz"},
-            ],
-            simulate_failure="flush",
+            simulate_send_failure=True,
             expected_completions=[],
             expected_state={"state": "interact", "write_failed": True},
         )