schema.py

"""Code generation script.

Builds the schema from the public endpoint
based on the OpenAPI specification.
"""

import os
import subprocess

import requests
from pydantic import Field
from pydantic_core import Url
from pydantic_settings import BaseSettings


class Settings(BaseSettings):
    """Settings for the code generation.

    Each setting can be overridden by an environment variable.

    Parameters
    ----------
    out_dir : str
        The directory where the generated schema will be saved.

    module_name : str
        The name of the module to be generated.

    openapi_source : str
        The source of the OpenAPI schema.

    schema_name : str
        The name of the schema file to be generated.

    delete_after : bool
        Whether to delete the schema file after the code generation.

    """

    out_dir: str = Field("qcog_python_client/schema/generated_schema/", alias="OUT_DIR")
    module_name: str = Field("models", alias="MODULE_NAME")
    openapi_source: Url = Field(
        "https://dev.qognitive.io/openapi.json", alias="OPENAPI_SOURCE"
    )
    schema_name: str = Field("schema", alias="SCHEMA_NAME")
    delete_after: bool = Field(True, alias="DELETE_AFTER")


def pull_openapi_schema(from_source: str, schemafile_name: str) -> str:
    """Pull the OpenAPI schema from the source.

    Generate a new schema.json file from the source at the given address.
    """
    response = requests.get(from_source, timeout=10)
    schema_path = os.path.join(os.getcwd(), schemafile_name)

    with open(schema_path, "w") as f:
        f.write(response.text)

    return schema_path


def main():
    """Generate the schema."""
    settings = Settings()

    # Pull the OPENAPI schema from the source

    schema_addr = pull_openapi_schema(
        from_source=settings.openapi_source,
        schemafile_name=(
            f"{settings.schema_name if settings.schema_name.endswith('.json') else settings.schema_name + '.json'}"  # noqa E501
        ),
    )

    # Build the paths for the codegen output
    codegen_out_dir = os.path.join(os.getcwd(), settings.out_dir)
    codegen_out_path = os.path.join(codegen_out_dir, settings.module_name)

    # Add extension if not present
    codegen_out_path = (
        codegen_out_path
        if codegen_out_path.endswith(".py")
        else codegen_out_path + ".py"
    )

    # Create the folder if it doesn't exist
    if not os.path.exists(codegen_out_dir):
        os.makedirs(os.path.dirname(codegen_out_dir), exist_ok=True)

    returncode = subprocess.call(
        [
            "datamodel-codegen",
            "--input",
            schema_addr,
            "--output",
            codegen_out_path,
            "--use-schema-description",
            "--output-model-type",
            "pydantic_v2.BaseModel",
            "--use-one-literal-as-default",  # Set a default value for literals
            "--allow-extra-fields",  # Allow new fields in the model. This is to avoid breaking changes if the API gets updates with new fields # noqa E501
            "--custom-file-header",
            "# mypy: ignore-errors\n# Autogenerated schema file\n# Tool: https://github.com/koxudaxi/datamodel-code-generator .",  # noqa E501
        ],
        timeout=10,
    )

    # Wait for the process to finish
    if returncode != 0:
        raise Exception("Failed to generate the schema")
    print("Schema generated successfully")

    # Delete the schema file if specified
    if settings.delete_after:
        os.remove(schema_addr)


if __name__ == "__main__":
    main()