22# Licensed under the MIT License.
33"""Lightweight JSON-RPC over standard IO."""
44
5+ from __future__ import annotations
6+
57import atexit
68import io
79import json
1012import threading
1113import uuid
1214from concurrent .futures import ThreadPoolExecutor
13- from typing import BinaryIO , Dict , Optional , Sequence , Union , cast
15+ from typing import BinaryIO , Sequence , cast
1416
1517CONTENT_LENGTH = "Content-Length: "
1618RUNNER_SCRIPT = str (pathlib .Path (__file__ ).parent / "runner.py" )
@@ -48,7 +50,9 @@ def write(self, data):
4850 with self ._lock :
4951 content = json .dumps (data )
5052 length = len (content .encode ("utf-8" ))
51- self ._writer .write (f"{ CONTENT_LENGTH } { length } \r \n \r \n { content } " .encode ("utf-8" ))
53+ self ._writer .write (
54+ f"{ CONTENT_LENGTH } { length } \r \n \r \n { content } " .encode () # type: ignore
55+ )
5256 self ._writer .flush ()
5357
5458
@@ -116,16 +120,18 @@ def receive_data(self):
116120
117121def create_json_rpc (readable : BinaryIO , writable : BinaryIO ) -> JsonRpc :
118122 """Creates JSON-RPC wrapper for the readable and writable streams."""
119- return JsonRpc (cast (io .TextIOWrapper , readable ), cast (io .TextIOWrapper , writable ))
123+ return JsonRpc (
124+ cast (io .TextIOWrapper , readable ), cast (io .TextIOWrapper , writable )
125+ )
120126
121127
122128class ProcessManager :
123129 """Manages sub-processes launched for running tools."""
124130
125131 def __init__ (self ):
126- self ._args : Dict [str , Sequence [str ]] = {}
127- self ._processes : Dict [str , subprocess .Popen ] = {}
128- self ._rpc : Dict [str , JsonRpc ] = {}
132+ self ._args : dict [str , Sequence [str ]] = {}
133+ self ._processes : dict [str , subprocess .Popen ] = {}
134+ self ._rpc : dict [str , JsonRpc ] = {}
129135 self ._lock = threading .Lock ()
130136 self ._thread_pool = ThreadPoolExecutor (10 )
131137
@@ -139,7 +145,9 @@ def stop_all_processes(self):
139145 pass
140146 self ._thread_pool .shutdown (wait = False )
141147
142- def start_process (self , workspace : str , args : Sequence [str ], cwd : str ) -> None :
148+ def start_process (
149+ self , workspace : str , args : Sequence [str ], cwd : str
150+ ) -> None :
143151 """Starts a process and establishes JSON-RPC communication over stdio."""
144152 proc = subprocess .Popen (
145153 args ,
@@ -175,7 +183,7 @@ def get_json_rpc(self, workspace: str) -> JsonRpc:
175183_process_manager = ProcessManager ()
176184
177185
178- def _get_json_rpc (workspace : str ) -> Union [ JsonRpc , None ] :
186+ def _get_json_rpc (workspace : str ) -> JsonRpc | None :
179187 try :
180188 return _process_manager .get_json_rpc (workspace )
181189 except StreamClosedException :
@@ -186,7 +194,7 @@ def _get_json_rpc(workspace: str) -> Union[JsonRpc, None]:
186194
187195def get_or_start_json_rpc (
188196 workspace : str , interpreter : Sequence [str ], cwd : str
189- ) -> Union [ JsonRpc , None ] :
197+ ) -> JsonRpc | None :
190198 """Gets an existing JSON-RPC connection or starts one and return it."""
191199 res = _get_json_rpc (workspace )
192200 if not res :
@@ -199,10 +207,10 @@ def get_or_start_json_rpc(
199207class RpcRunResult :
200208 """Object to hold result from running tool over RPC."""
201209
202- def __init__ (self , stdout : str , stderr : str , exception : Optional [ str ] = None ):
210+ def __init__ (self , stdout : str , stderr : str , exception : str | None = None ):
203211 self .stdout : str = stdout
204212 self .stderr : str = stderr
205- self .exception : Optional [ str ] = exception
213+ self .exception : str | None = exception
206214
207215
208216def run_over_json_rpc (
@@ -212,10 +220,10 @@ def run_over_json_rpc(
212220 argv : Sequence [str ],
213221 use_stdin : bool ,
214222 cwd : str ,
215- source : str = None ,
223+ source : str | None = None ,
216224) -> RpcRunResult :
217225 """Uses JSON-RPC to execute a command."""
218- rpc : Union [ JsonRpc , None ] = get_or_start_json_rpc (workspace , interpreter , cwd )
226+ rpc : JsonRpc | None = get_or_start_json_rpc (workspace , interpreter , cwd )
219227 if not rpc :
220228 raise Exception ("Failed to run over JSON-RPC." )
221229
@@ -236,7 +244,9 @@ def run_over_json_rpc(
236244 data = rpc .receive_data ()
237245
238246 if data ["id" ] != msg_id :
239- return RpcRunResult ("" , f"Invalid result for request: { json .dumps (msg , indent = 4 )} " )
247+ return RpcRunResult (
248+ "" , f"Invalid result for request: { json .dumps (msg , indent = 4 )} "
249+ )
240250
241251 if "error" in data :
242252 if data .get ("exception" , False ):
0 commit comments