client.py

from rich.console import Console
from rich.markdown import Markdown
import os

import grpc
from mpai_cae_arp.network import arp_pb2
from mpai_cae_arp.network import arp_pb2_grpc

channels = {
    "AudioAnalyser": grpc.insecure_channel("[::]:50051"),
    "VideoAnalyser": grpc.insecure_channel("[::]:50052"),
    "TapeIrregularityClassifier": grpc.insecure_channel("[::]:50051/tape-irregularity-classifier"),
    "TapeAudioRestoration": grpc.insecure_channel("[::]:50051/tape-audio-restoration"),
    "Packager": grpc.insecure_channel("[::]:50051/packager"),
}

def run(console: Console):

    audio_analyser = arp_pb2_grpc.AIMStub(channels["AudioAnalyser"])
    video_analyser = arp_pb2_grpc.AIMStub(channels["VideoAnalyser"])

    request = arp_pb2.InfoRequest()
    for analyser in [audio_analyser, video_analyser]:
        response = analyser.getInfo(request)
        console.print("[bold]{}[/], v{}".format(response.title, response.version))
        console.print(Markdown(response.description))

    request = arp_pb2.JobRequest(
        working_dir="../data",
        files_name="BERIO100",
        index=1,
    )

    with console.status("[bold]Computing AudioAnalyser IrregularityFile 1...", spinner="bouncingBall"):
        for result in audio_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    request.files_name = "BERIO100.mov"
    with console.status("[bold]Computing VideoAnalyser IrregularityFiles...", spinner="bouncingBall"):
        for result in video_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    request.index = 2
    request.files_name = "BERIO100"
    with console.status("[bold]Computing AudioAnalyser IrregularityFile 2...", spinner="bouncingBall"):
        for result in audio_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    channels["AudioAnalyser"].close()
    channels["VideoAnalyser"].close()

    console.print("[bold green]Success![/] :tada:")


if __name__ == '__main__':
    console = Console()
    run(console)