4040__title__ = "Mypy Daemon Client"
4141__license__ = "MIT"
4242
43- import base64
4443import contextlib
44+ import struct
4545import sys
4646from collections import ChainMap
47- from typing import TYPE_CHECKING , Literal , TypedDict , cast
47+ from typing import TYPE_CHECKING , Final , Literal , TypedDict , cast
4848
4949import orjson
5050
5959 daemonize as _daemonize ,
6060 process_start_options as _process_start_options ,
6161 )
62- from mypy .ipc import IPCClient as _IPCClient , IPCException as _IPCException
62+ from mypy .ipc import (
63+ HEADER_SIZE ,
64+ IPCClient as _IPCClient ,
65+ IPCException as _IPCException ,
66+ )
6367 from mypy .version import __version__
6468
6569if TYPE_CHECKING :
@@ -101,6 +105,14 @@ class Stats(TypedDict):
101105 sccs_left : int
102106 nodes_left : int
103107 cache_commit_time : float
108+ type_expression_parse_count : int
109+ type_expression_full_parse_success_count : int
110+ type_expression_full_parse_failure_count : int
111+ load_missing_time : float
112+ order_scc_time : float
113+ semanal_time : float
114+ type_check_time : float
115+ flush_and_cache_time : float
104116
105117
106118class Response (TypedDict ):
@@ -256,6 +268,11 @@ def _read_request_response_json(request_response: str | bytes) -> Response:
256268 return cast ("Response" , data )
257269
258270
271+ IPC_CONNECTION_CLOSED_ERROR : Final = Response (
272+ {"error" : "No data received, IPC socket connection closed." },
273+ )
274+
275+
259276async def _request_win32 (
260277 name : str ,
261278 request_arguments : bytes ,
@@ -273,7 +290,7 @@ async def _receive(
273290 """
274291 bdata : str = await async_connection .read ()
275292 if not bdata :
276- return { "error" : "No data received, IPC socket connection closed." }
293+ return IPC_CONNECTION_CLOSED_ERROR
277294 return _read_request_response_json (bdata )
278295
279296 try :
@@ -295,19 +312,24 @@ async def _receive(
295312 return {"error" : str (err )}
296313
297314
298- def find_frame_in_buffer (
315+ def frame_from_buffer (
299316 buffer : bytearray ,
300- ) -> tuple [bytearray , bytearray | None ]:
317+ message_size : int | None = None ,
318+ ) -> tuple [bytearray , bytes | None , int | None ]:
301319 """Return a full frame from the bytes we have in the buffer.
302320
303- Returns (next frame keep data, complete frame | None)
321+ Returns (next frame keep data, complete frame | None, complete frame size | None )
304322 """
305- space_pos = buffer .find (b" " )
306- if space_pos == - 1 :
307- # Incomplete frame
308- return buffer , None
309- # We have a full frame
310- return buffer [space_pos + 1 :], buffer [:space_pos ]
323+ size = len (buffer )
324+ if size < HEADER_SIZE :
325+ return buffer , None , None
326+ if message_size is None :
327+ message_size = struct .unpack ("!L" , buffer [:HEADER_SIZE ])[0 ]
328+ if size < message_size + HEADER_SIZE :
329+ return buffer , None , message_size
330+ # We have a full frame, avoid extra copy in case we get a large frame.
331+ bdata = memoryview (buffer )[HEADER_SIZE : HEADER_SIZE + message_size ]
332+ return buffer [HEADER_SIZE + message_size :], bytes (bdata ), message_size
311333
312334
313335async def _request_linux (
@@ -317,38 +339,56 @@ async def _request_linux(
317339) -> Response :
318340 """Request from daemon on linux/unix."""
319341 buffer = bytearray ()
320- frame : bytearray | None = None
342+ frame : bytes | None = None
321343 all_responses : list [Response ] = []
322344 async with await trio .open_unix_socket (filename ) as connection :
323- # Frame the data by urlencoding it and separating by space.
324- request_frame = base64 .encodebytes (request_arguments ) + b" "
345+ # Frame the request and send.
346+ request_frame = (
347+ struct .pack ("!L" , len (request_arguments )) + request_arguments
348+ )
325349 await connection .send_all (request_frame )
326350
327351 is_not_done = True
328352 while is_not_done :
329- # Receive more data into the buffer.
330- try :
331- if timeout is None :
332- more = await connection .receive_some ()
333- else :
334- with trio .fail_after (timeout ):
335- more = await connection .receive_some ()
336- except trio .TooSlowError :
337- return {"error" : "IPC socket connection timed out" }
338- if not more :
339- # Connection closed
340- # Socket was empty and we didn't get any frame.
341- return {
342- "error" : "No data received, IPC socket connection closed." ,
343- }
344- buffer .extend (more )
345-
346- buffer , frame = find_frame_in_buffer (buffer )
353+ # check if we have read entire frame yet
354+ buffer , frame , complete_frame_size = frame_from_buffer (buffer )
355+
356+ max_read_bytes : int | None
357+ if complete_frame_size is None :
358+ # if have not read frame size header, do so.
359+ max_read_bytes = HEADER_SIZE
360+ else :
361+ # number of bytes to read next is the complete frame
362+ # size - the size header - current bytes count
363+ max_read_bytes = max (
364+ HEADER_SIZE ,
365+ complete_frame_size - len (buffer ) + HEADER_SIZE ,
366+ )
367+
347368 if frame is None :
369+ # Have not read entire frame yet, so receive more data
370+ # into the buffer.
371+ try :
372+ if timeout is None :
373+ more = await connection .receive_some (max_read_bytes )
374+ else :
375+ with trio .fail_after (timeout ):
376+ more = await connection .receive_some (
377+ max_read_bytes ,
378+ )
379+ except trio .TooSlowError :
380+ return {
381+ "error" : f"IPC socket connection timed out ({ timeout = } )" ,
382+ }
383+ if not more :
384+ # Connection closed
385+ # Socket was empty and we didn't get any frame.
386+ return IPC_CONNECTION_CLOSED_ERROR
387+ buffer .extend (more )
348388 continue
389+
349390 # Frame is not None, we read a full frame
350- response_text = base64 .decodebytes (frame )
351- response = _read_request_response_json (response_text )
391+ response = _read_request_response_json (frame )
352392
353393 is_not_done = not bool (response .pop ("final" , False ))
354394 all_responses .append (response )
0 commit comments