pydantic_glue/cli.py

"""CLI to convert pydantic class to AWS Glue schema."""

import json
import logging
import sys
from argparse import ArgumentParser
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

from pydantic_glue import convert

logging.basicConfig(format="{asctime} - {levelname} - {message}", level=logging.INFO, style="{", stream=sys.stdout)
log = logging.getLogger(__name__)


@dataclass
class Arguments:
    """CLI Arguments."""

    module_file: str
    class_name: str
    output_file: str
    log_result: bool
    json_schema_by_alias: bool


def parse_args(argv: list[str]) -> Arguments:
    """Parse CLI arguments."""
    parser = ArgumentParser()
    parser.add_argument("-f", dest="source_file", required=True, type=str, help="Path to the python file")
    parser.add_argument("-c", dest="class_name", required=True, type=str, help="Python class name")
    parser.add_argument("-o", dest="output_file", required=False, type=str, help="Path to the output file")
    parser.add_argument(
        "-l", dest="log_result", action="store_true", default=False, help="Flag if need to print result in log"
    )
    parser.add_argument(
        "--schema-by-name",
        dest="json_schema_by_alias",
        action="store_false",
        default=True,
        help="Flag to not use name for json schema generation",
    )
    args = parser.parse_args(argv)

    return Arguments(
        module_file=args.source_file.removesuffix(".py"),
        class_name=args.class_name,
        output_file=args.output_file,
        log_result=args.log_result,
        json_schema_by_alias=args.json_schema_by_alias,
    )


def cli() -> None:
    """CLI entry point."""
    args = parse_args(sys.argv[1:])
    sys.path.append(Path(args.module_file).parent.as_posix())
    imported = __import__(Path(args.module_file).name)

    model = getattr(imported, args.class_name)
    input_schema = json.dumps(model.model_json_schema(by_alias=args.json_schema_by_alias))
    converted = convert(input_schema)

    output = json.dumps(
        {
            "//": f"Generated by pydantic-glue at {datetime.now(tz=timezone.utc)}. DO NOT EDIT",
            "columns": dict(converted),
        },
        indent=2,
    )

    if args.log_result or not args.output_file:
        log.info(f"Generated file content:\n{output}")

    if args.output_file:
        ouf = Path(args.output_file)
        ouf.parent.mkdir(parents=True, exist_ok=True)
        ouf.write_text(output)