tts_service/cli.py

import os
from pathlib import Path

import boto3
import click
from click_help_colors import HelpColorsGroup
from dotenv import load_dotenv

load_dotenv()


@click.group(
    cls=HelpColorsGroup,
    help_headers_color="yellow",
    help_options_color="green",
)
def main() -> None:
    """TTS Service CLI"""


@main.command()
@click.option("--share", is_flag=True, help="Share the service")
def serve(share: bool) -> None:
    """Start the TTS Service"""
    from tts_service.app import app

    app.launch(share=share)


@main.group()
def service() -> None:
    """Manages the deployed service."""


@service.command()
@click.option("--bucket", "-b", default=lambda: os.environ["BUCKET"], help="the bucket to upload voices to")
@click.option("--prefix", "-p", default=lambda: os.environ["VOICES_KEY_PREFIX"], help="the prefix to use for the keys")
@click.option("--delete", is_flag=True, help="delete extraneous files from dest")
@click.option("--dry-run", "-n", is_flag=True, help="perform a trial run with no changes made")
@click.argument("directory", type=click.Path(exists=True, file_okay=False, path_type=Path), nargs=1)
def upload_voices(bucket: str, prefix: str, delete: bool, dry_run: bool, directory: Path) -> None:
    """Upload voices to the service"""
    s3 = boto3.client("s3")
    prefix = prefix.strip("/")
    names = set()
    for path in directory.glob("*.pth"):
        names.add(path.name)
        with path.open("rb") as file:
            if dry_run:
                click.echo(f"Would upload {path.name} to {bucket}/{prefix}")
            else:
                s3.put_object(Bucket=bucket, Key=f"{prefix}/{path.name}", Body=file)
                # s3.upload_fileobj(file, bucket, f"{prefix}/{path.name}")
    if not names:
        raise click.ClickException(f"no voices found in directory {directory}")
    deleted = 0
    if delete:
        paginator = s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page["Contents"]:
                key = obj["Key"]
                if key.split("/")[-1] not in names:
                    if dry_run:
                        click.echo(f"Would delete {key}")
                    else:
                        s3.delete_object(Bucket=bucket, Key=key)
                    deleted += 1
    deleted_message = f", {deleted} deleted" if delete else ""
    if not dry_run:
        click.echo(f"{bucket}/{prefix}: {len(names)} voices uploaded{deleted_message}")


if __name__ == "__main__":
    main()