Stream dmypy output instead of dumping everything at the end

This does 2 things:
1. It changes the IPC code to work with multiple messages.
2. It changes the dmypy client/server communication so that it streams
stdout/stderr instead of dumping everything at the end.

For 1, we have to provide a way to separate out different messages.
We can frame messages as bytes separated by whitespace character. That
means we have to encode the message in a scheme that escapes whitespace.
The urllib.parse quote/unquote seems reasonable. It encodes more than
needed but the application is not IPC IO limited so it should be fine.
With this convention in place, all we have to do is read from the socket
stream until we have a whitespace character.

For 2, since we communicate with JSONs, it's easy to add a "finished"
key that tells us it's the final response from dmypy. Anything else is
just stdout/stderr output.

Note: dmypy server also returns out/err which is the output of actual
mypy type checking. Right now this change does not stream that output.
We can stream that in a followup change. We just have to decide on how
to differenciate the 4 text streams (stdout/stderr/out/err) that will
now be interleaved.

Played around with it on Linux quite a bit. Will also test it some more
on Windows.

The WriteToConn class could use more love. I just put a bare minimum to
test the rest.

Author: svalentin

mypy/dmypy/client.py

@@ -17,7 +17,7 @@
 from typing import Any, Callable, Mapping, NoReturn
 
 from mypy.dmypy_os import alive, kill
-from mypy.dmypy_util import DEFAULT_STATUS_FILE, receive
+from mypy.dmypy_util import DEFAULT_STATUS_FILE, receive, send
 from mypy.ipc import IPCClient, IPCException
 from mypy.util import check_python_version, get_terminal_width, should_force_color
 from mypy.version import __version__
@@ -659,28 +659,29 @@ def request(
     # so that it can format the type checking output accordingly.
     args["is_tty"] = sys.stdout.isatty() or should_force_color()
     args["terminal_width"] = get_terminal_width()
-    bdata = json.dumps(args).encode("utf8")
     _, name = get_status(status_file)
     try:
         with IPCClient(name, timeout) as client:
-            client.write(bdata)
-            response = receive(client)
+            send(client, args)
+
+            final = False
+            while not final:
+                response = receive(client)
+                final = bool(response.pop("final", False))
+                # Display debugging output written to stdout/stderr in the server process for convenience.
+                # This should not be confused with "out" and "err" fields in the response.
+                # Those fields hold the output of the "check" command, and are handled in check_output().
+                stdout = response.pop("stdout", None)
+                if stdout:
+                    sys.stdout.write(stdout)
+                stderr = response.pop("stderr", None)
+                if stderr:
+                    sys.stderr.write(stderr)
     except (OSError, IPCException) as err:
         return {"error": str(err)}
     # TODO: Other errors, e.g. ValueError, UnicodeError
-    else:
-        # Display debugging output written to stdout/stderr in the server process for convenience.
-        # This should not be confused with "out" and "err" fields in the response.
-        # Those fields hold the output of the "check" command, and are handled in check_output().
-        stdout = response.get("stdout")
-        if stdout:
-            sys.stdout.write(stdout)
-        stderr = response.get("stderr")
-        if stderr:
-            print("-" * 79)
-            print("stderr:")
-            sys.stdout.write(stderr)
-        return response
+
+    return response
 
 
 def get_status(status_file: str) -> tuple[int, str]:

mypy/dmypy_server.py

@@ -17,18 +17,18 @@
 import time
 import traceback
 from contextlib import redirect_stderr, redirect_stdout
-from typing import AbstractSet, Any, Callable, Final, List, Sequence, Tuple
+from typing import AbstractSet, Any, Callable, Final, Iterable, List, Sequence, Tuple
 from typing_extensions import TypeAlias as _TypeAlias
 
 import mypy.build
 import mypy.errors
 import mypy.main
-from mypy.dmypy_util import receive
+from mypy.dmypy_util import receive, send
 from mypy.find_sources import InvalidSourceList, create_source_list
 from mypy.fscache import FileSystemCache
 from mypy.fswatcher import FileData, FileSystemWatcher
 from mypy.inspections import InspectionEngine
-from mypy.ipc import IPCServer
+from mypy.ipc import IPCBase, IPCServer
 from mypy.modulefinder import BuildSource, FindModuleCache, SearchPaths, compute_search_paths
 from mypy.options import Options
 from mypy.server.update import FineGrainedBuildManager, refresh_suppressed_submodules
@@ -208,19 +208,34 @@ def _response_metadata(self) -> dict[str, str]:
 
     def serve(self) -> None:
         """Serve requests, synchronously (no thread or fork)."""
+
+        class WriteToConn(object):
+            def __init__(self, server: IPCBase, output_key: str = "stdout"):
+                self.server = server
+                self.output_key = output_key
+
+            def write(self, output: str) -> int:
+                resp: dict[str, Any] = {}
+                resp[self.output_key] = output
+                send(server, resp)
+                return len(output)
+
+            def writelines(self, lines: Iterable[str]) -> None:
+                for s in lines:
+                    self.write(s)
+
         command = None
         server = IPCServer(CONNECTION_NAME, self.timeout)
+
         try:
             with open(self.status_file, "w") as f:
                 json.dump({"pid": os.getpid(), "connection_name": server.connection_name}, f)
                 f.write("\n")  # I like my JSON with a trailing newline
             while True:
                 with server:
                     data = receive(server)
-                    debug_stdout = io.StringIO()
-                    debug_stderr = io.StringIO()
-                    sys.stdout = debug_stdout
-                    sys.stderr = debug_stderr
+                    sys.stdout = WriteToConn(server, "stdout")  # type: ignore[assignment]
+                    sys.stderr = WriteToConn(server, "stderr")  # type: ignore[assignment]
                     resp: dict[str, Any] = {}
                     if "command" not in data:
                         resp = {"error": "No command found in request"}
@@ -237,15 +252,13 @@ def serve(self) -> None:
                                 tb = traceback.format_exception(*sys.exc_info())
                                 resp = {"error": "Daemon crashed!\n" + "".join(tb)}
                                 resp.update(self._response_metadata())
-                                resp["stdout"] = debug_stdout.getvalue()
-                                resp["stderr"] = debug_stderr.getvalue()
-                                server.write(json.dumps(resp).encode("utf8"))
+                                resp["final"] = True
+                                send(server, resp)
                                 raise
-                    resp["stdout"] = debug_stdout.getvalue()
-                    resp["stderr"] = debug_stderr.getvalue()
+                    resp["final"] = True
                     try:
                         resp.update(self._response_metadata())
-                        server.write(json.dumps(resp).encode("utf8"))
+                        send(server, resp)
                     except OSError:
                         pass  # Maybe the client hung up
                     if command == "stop":

mypy/dmypy_util.py

@@ -14,7 +14,7 @@
 
 
 def receive(connection: IPCBase) -> Any:
-    """Receive JSON data from a connection until EOF.
+    """Receive single JSON data frame from a connection.
 
     Raise OSError if the data received is not valid JSON or if it is
     not a dict.
@@ -23,9 +23,19 @@ def receive(connection: IPCBase) -> Any:
     if not bdata:
         raise OSError("No data received")
     try:
-        data = json.loads(bdata.decode("utf8"))
+        data = json.loads(bdata)
     except Exception as e:
         raise OSError("Data received is not valid JSON") from e
     if not isinstance(data, dict):
         raise OSError(f"Data received is not a dict ({type(data)})")
     return data
+
+def send(connection: IPCBase, data: Any) -> None:
+    """Send data to a connection encoded and framed.
+
+    The data must be JSON-serializable. We assume that a single send call is a
+    single frame to be sent on the connect.
+    As an easy way to separate frames, we urlencode them and separate by space.
+    Last frame also has a trailing space.
+    """
+    connection.write(json.dumps(data))

mypy/ipc.py

@@ -13,6 +13,7 @@
 import tempfile
 from types import TracebackType
 from typing import Callable, Final
+from urllib.parse import quote,unquote
 
 if sys.platform == "win32":
     # This may be private, but it is needed for IPC on Windows, and is basically stable
@@ -47,8 +48,9 @@ class IPCBase:
     def __init__(self, name: str, timeout: float | None) -> None:
         self.name = name
         self.timeout = timeout
+        self.buffer = bytearray()
 
-    def read(self, size: int = 100000) -> bytes:
+    def read(self, size: int = 100000) -> str:
         """Read bytes from an IPC connection until its empty."""
         bdata = bytearray()
         if sys.platform == "win32":
@@ -66,7 +68,13 @@ def read(self, size: int = 100000) -> bytes:
                 _, err = ov.GetOverlappedResult(True)
                 more = ov.getbuffer()
                 if more:
-                    bdata.extend(more)
+                    self.buffer.extend(more)
+                    space_pos = self.buffer.find(b" ")
+                    if space_pos != -1:
+                        # We have a full frame
+                        bdata = self.buffer[: space_pos]
+                        self.buffer = self.buffer[space_pos + 1 :]
+                        break
                 if err == 0:
                     # we are done!
                     break
@@ -79,15 +87,26 @@ def read(self, size: int = 100000) -> bytes:
             while True:
                 more = self.connection.recv(size)
                 if not more:
+                    # Connection closed
                     break
-                bdata.extend(more)
-        return bytes(bdata)
+                self.buffer.extend(more)
+                space_pos = self.buffer.find(b" ")
+                if space_pos != -1:
+                    # We have a full frame
+                    bdata = self.buffer[: space_pos]
+                    self.buffer = self.buffer[space_pos + 1 :]
+                    break
+        return unquote(bytes(bdata).decode("utf8"))
 
-    def write(self, data: bytes) -> None:
+    def write(self, data: str) -> None:
         """Write bytes to an IPC connection."""
+
+        # Frame the data by urlencoding it and separating by space.
+        encoded_data = (quote(data) + " ").encode("utf8")
+
         if sys.platform == "win32":
             try:
-                ov, err = _winapi.WriteFile(self.connection, data, overlapped=True)
+                ov, err = _winapi.WriteFile(self.connection, encoded_data, overlapped=True)
                 try:
                     if err == _winapi.ERROR_IO_PENDING:
                         timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
@@ -101,12 +120,11 @@ def write(self, data: bytes) -> None:
                     raise
                 bytes_written, err = ov.GetOverlappedResult(True)
                 assert err == 0, err
-                assert bytes_written == len(data)
+                assert bytes_written == len(encoded_data)
             except OSError as e:
                 raise IPCException(f"Failed to write with error: {e.winerror}") from e
         else:
-            self.connection.sendall(data)
-            self.connection.shutdown(socket.SHUT_WR)
+            self.connection.sendall(encoded_data)
 
     def close(self) -> None:
         if sys.platform == "win32":

mypy/test/testipc.py

@@ -15,13 +15,23 @@
 def server(msg: str, q: Queue[str]) -> None:
     server = IPCServer(CONNECTION_NAME)
     q.put(server.connection_name)
-    data = b""
+    data = ""
     while not data:
         with server:
-            server.write(msg.encode())
+            server.write(msg)
             data = server.read()
     server.cleanup()
 
+def server_multi_message_echo(q: Queue[str]) -> None:
+    server = IPCServer(CONNECTION_NAME)
+    q.put(server.connection_name)
+    data = ""
+    with server:
+        while data != "quit":
+            data = server.read()
+            server.write(data)
+    server.cleanup()
+
 
 class IPCTests(TestCase):
     def test_transaction_large(self) -> None:
@@ -31,8 +41,8 @@ def test_transaction_large(self) -> None:
         p.start()
         connection_name = queue.get()
         with IPCClient(connection_name, timeout=1) as client:
-            assert client.read() == msg.encode()
-            client.write(b"test")
+            assert client.read() == msg
+            client.write("test")
         queue.close()
         queue.join_thread()
         p.join()
@@ -44,12 +54,29 @@ def test_connect_twice(self) -> None:
         p.start()
         connection_name = queue.get()
         with IPCClient(connection_name, timeout=1) as client:
-            assert client.read() == msg.encode()
-            client.write(b"")  # don't let the server hang up yet, we want to connect again.
+            assert client.read() == msg
+            client.write("")  # don't let the server hang up yet, we want to connect again.
+
+        with IPCClient(connection_name, timeout=1) as client:
+            assert client.read() == msg
+            client.write("test")
+        queue.close()
+        queue.join_thread()
+        p.join()
+        assert p.exitcode == 0
 
+    def test_multiple_messages(self) -> None:
+        queue: Queue[str] = Queue()
+        p = Process(target=server_multi_message_echo, args=(queue,), daemon=True)
+        p.start()
+        connection_name = queue.get()
         with IPCClient(connection_name, timeout=1) as client:
-            assert client.read() == msg.encode()
-            client.write(b"test")
+            client.write("test1")
+            assert client.read() == "test1"
+            client.write("test2")
+            assert client.read() == "test2"
+            client.write("quit")
+            assert client.read() == "quit"
         queue.close()
         queue.join_thread()
         p.join()