11from __future__ import annotations
22
3+ import logging
34import socket
45from pathlib import Path
56from typing import Annotated , Any , Literal
1213 alias_generators ,
1314)
1415
16+ logger = logging .getLogger (__name__ )
17+
1518
1619class ProcessorAccess :
17- def __init__ (self , hostname : str , port : int ):
20+ def __init__ (
21+ self ,
22+ hostname : str ,
23+ port : int ,
24+ * ,
25+ log_level : int = logging .DEBUG ,
26+ ):
1827 self .hostname = hostname
1928 self .port = port
29+ self .log_level = log_level
2030 self .socket = socket .socket (socket .AF_INET , socket .SOCK_STREAM )
2131
2232 def connect (self ):
@@ -31,32 +41,47 @@ def flash(self, path: str | Path):
3141 raise ValueError ("Path must be absolute." )
3242
3343 self ._send_request (FlashRequest (path = path ))
34- return self ._recv_response ()
35-
36- def dump (self , path : str | Path ):
44+ return self ._recv_response (SuccessResponse )
45+
46+ def dump (
47+ self ,
48+ path : str | Path ,
49+ address : int | None = None ,
50+ bytes : int | None = None ,
51+ ):
3752 path = Path (path )
3853 if not path .is_absolute ():
3954 raise ValueError ("Path must be absolute." )
4055
41- self ._send_request (DumpRequest (path = path ))
42- return self ._recv_response ()
56+ self ._send_request (DumpRequest (path = path , address = address , bytes = bytes ))
57+ return self ._recv_response (SuccessResponse )
4358
4459 def start (self , wait : bool ):
4560 self ._send_request (StartRequest (wait = wait ))
46- return self ._recv_response ()
61+ return self ._recv_response (SuccessResponse )
4762
4863 def stop (self ):
4964 self ._send_request (StopRequest ())
50- return self ._recv_response ()
65+ return self ._recv_response (SuccessResponse )
66+
67+ def status (self ):
68+ self ._send_request (StatusRequest ())
69+ return self ._recv_response (StatusResponse )
5170
5271 def _send_request (self , request : Request ) -> None :
5372 message = request .model_dump_json () + "\n "
73+ logger .log (self .log_level , f"Sending request: { message .rstrip ()} " )
5474 self .socket .sendall (message .encode ("utf-8" ))
5575
56- def _recv_response (self ) -> Response :
76+ def _recv_response [ T ] (self , response_type : type [ T ] ) -> T :
5777 with self .socket .makefile ("r" , encoding = "utf-8" ) as f :
5878 line = f .readline ()
59- return response_ta .validate_json (line )
79+ logger .log (self .log_level , f"Received response: { line .rstrip ()} " )
80+ match TypeAdapter (response_type | ErrorResponse ).validate_json (line ):
81+ case ErrorResponse () as e :
82+ raise ProcessorError (e )
83+ case response :
84+ return response
6085
6186 def __enter__ (self ):
6287 self .connect ()
@@ -81,6 +106,8 @@ class FlashRequest(BaseModel):
81106class DumpRequest (BaseModel ):
82107 type : Literal ["dump" ] = "dump"
83108 path : Path
109+ address : int | None
110+ bytes : int | None
84111
85112
86113class StartRequest (BaseModel ):
@@ -92,8 +119,12 @@ class StopRequest(BaseModel):
92119 type : Literal ["stop" ] = "stop"
93120
94121
122+ class StatusRequest (BaseModel ):
123+ type : Literal ["status" ] = "status"
124+
125+
95126type Request = Annotated [
96- FlashRequest | DumpRequest | StartRequest | StopRequest ,
127+ FlashRequest | DumpRequest | StartRequest | StopRequest | StatusRequest ,
97128 Field (discriminator = "type" ),
98129]
99130
@@ -103,14 +134,25 @@ class SuccessResponse(BaseModel):
103134 message : str
104135
105136
137+ class StatusResponse (BaseModel ):
138+ type : Literal ["status" ]
139+ running : bool
140+ pc : int | None
141+ error_output : str
142+
143+
106144class ErrorResponse (BaseModel ):
107145 type : Literal ["error" ]
108146 message : str
109147
110148
111149type Response = Annotated [
112- SuccessResponse | ErrorResponse ,
150+ SuccessResponse | StatusResponse | ErrorResponse ,
113151 Field (discriminator = "type" ),
114152]
115153
116- response_ta = TypeAdapter [Response ](Response )
154+
155+ class ProcessorError (RuntimeError ):
156+ def __init__ (self , response : ErrorResponse ):
157+ super ().__init__ (response .message )
158+ self .response = response
0 commit comments