async def stream_video(self, output_file: Optional[str], fps: Optional[int], format: VideoFormat) -> AsyncGenerator[bytes, None]: self.logger.info("Starting connection to backend") async with self.stub.video_stream.open() as stream: if self.is_local and output_file: self.logger.info( f"Streaming locally with companion writing to {output_file}" ) await stream.send_message( VideoStreamRequest(start=VideoStreamRequest.Start( file_path=output_file, fps=fps, format=VIDEO_FORMAT_MAP[format], ))) else: self.logger.info("Starting streaming over the wire") await stream.send_message( VideoStreamRequest(start=VideoStreamRequest.Start( file_path=None, fps=fps, format=VIDEO_FORMAT_MAP[format]))) try: iterator = generate_bytes(stream=stream, logger=self.logger) if output_file and not self.is_local: self.logger.info(f"Writing wired bytes to {output_file}") await drain_to_file(stream=iterator, file_path=output_file) else: async for data in iterator: yield data finally: self.logger.info("Stopping video streaming") await stream.send_message( VideoStreamRequest(stop=VideoStreamRequest.Stop())) await stream.end()
async def pull(self, bundle_id: str, src_path: str, dest_path: str) -> None: async with self.stub.pull.open() as stream: request = request = PullRequest( bundle_id=bundle_id, src_path=src_path, # not sending the destination to remote companion # so it streams the file back dst_path=dest_path if self.is_local else None, ) await stream.send_message(request) await stream.end() if self.is_local: await stream.recv_message() else: await drain_untar(generate_bytes(stream), output_path=dest_path) self.logger.info(f"pulled file to {dest_path}")
async def daemon(client: CompanionClient, request: PullRequest) -> PullResponse: destination = request.dst_path async with client.stub.pull.open() as stream: if not client.is_local: # not sending the destination to remote companion # so it streams the file back request = PullRequest(bundle_id=request.bundle_id, src_path=request.src_path, dst_path=None) await stream.send_message(request, end=True) if client.is_local: await stream.recv_message() else: await drain_untar(generate_bytes(stream), output_path=destination) client.logger.info(f"pulled file to {destination}") return PullResponse(payload=Payload(file_path=destination))