22# 
33# SPDX-License-Identifier: MIT 
44
5+ import  base64 
56from  pathlib  import  Path 
67from  typing  import  Any , Optional , Union 
78
89from  wokwi_client .framebuffer  import  (
9-     compare_framebuffer_png ,
10-     framebuffer_png_bytes ,
11-     framebuffer_read ,
10+     read_framebuffer_png_bytes ,
1211    save_framebuffer_png ,
1312)
1413
1514from  .__version__  import  get_version 
1615from  .constants  import  DEFAULT_WS_URL 
1716from  .control  import  set_control 
1817from  .event_queue  import  EventQueue 
19- from  .file_ops  import  download , download_file ,  upload , upload_file 
18+ from  .file_ops  import  download , upload , upload_file 
2019from  .pins  import  gpio_list , pin_listen , pin_read 
2120from  .protocol_types  import  EventMessage , ResponseMessage 
2221from  .serial  import  monitor_lines , write_serial 
@@ -93,17 +92,18 @@ async def upload_file(
9392        """ 
9493        return  await  upload_file (self ._transport , filename , local_path )
9594
96-     async  def  download (self , name : str ) ->  ResponseMessage :
95+     async  def  download (self , name : str ) ->  bytes :
9796        """ 
9897        Download a file from the simulator. 
9998
10099        Args: 
101100            name: The name of the file to download. 
102101
103102        Returns: 
104-             The response message from the server . 
103+             The downloaded file content as bytes . 
105104        """ 
106-         return  await  download (self ._transport , name )
105+         result  =  await  download (self ._transport , name )
106+         return  base64 .b64decode (result ["result" ]["binary" ])
107107
108108    async  def  download_file (self , name : str , local_path : Optional [Path ] =  None ) ->  None :
109109        """ 
@@ -113,7 +113,12 @@ async def download_file(self, name: str, local_path: Optional[Path] = None) -> N
113113            name: The name of the file to download. 
114114            local_path: The local path to save the downloaded file. If not provided, uses the name as the path. 
115115        """ 
116-         await  download_file (self ._transport , name , local_path )
116+         if  local_path  is  None :
117+             local_path  =  Path (name )
118+ 
119+         result  =  await  self .download (name )
120+         with  open (local_path , "wb" ) as  f :
121+             f .write (result )
117122
118123    async  def  start_simulation (
119124        self ,
@@ -267,22 +272,10 @@ async def set_control(
267272        """ 
268273        return  await  set_control (self ._transport , part = part , control = control , value = value )
269274
270-     async  def  framebuffer_read (self , id : str ) ->  ResponseMessage :
271-         """Read the current framebuffer for the given device id.""" 
272-         return  await  framebuffer_read (self ._transport , id = id )
273- 
274-     async  def  framebuffer_png_bytes (self , id : str ) ->  bytes :
275+     async  def  read_framebuffer_png_bytes (self , id : str ) ->  bytes :
275276        """Return the current framebuffer as PNG bytes.""" 
276-         return  await  framebuffer_png_bytes (self ._transport , id = id )
277+         return  await  read_framebuffer_png_bytes (self ._transport , id = id )
277278
278279    async  def  save_framebuffer_png (self , id : str , path : Path , overwrite : bool  =  True ) ->  Path :
279280        """Save the current framebuffer as a PNG file.""" 
280281        return  await  save_framebuffer_png (self ._transport , id = id , path = path , overwrite = overwrite )
281- 
282-     async  def  compare_framebuffer_png (
283-         self , id : str , reference : Path , save_mismatch : Optional [Path ] =  None 
284-     ) ->  bool :
285-         """Compare the current framebuffer with a reference PNG file.""" 
286-         return  await  compare_framebuffer_png (
287-             self ._transport , id = id , reference = reference , save_mismatch = save_mismatch 
288-         )
0 commit comments