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))
