cli/openbb_cli/argparse_translator/argparse_argument.py

"""Pydantic models for argparse arguments and argument groups."""

from typing import (
    Any,
    List,
    Literal,
    Optional,
    Tuple,
)

from pydantic import BaseModel, model_validator

SEP = "__"


class ArgparseArgumentModel(BaseModel):
    """Pydantic model for an argparse argument."""

    name: str
    type: Any
    dest: str
    default: Any
    required: bool
    action: Literal["store_true", "store"]
    help: Optional[str]
    nargs: Optional[Literal["+"]]
    choices: Optional[Tuple]

    @model_validator(mode="after")  # type: ignore
    @classmethod
    def validate_action(cls, values: "ArgparseArgumentModel"):
        """Validate the action based on the type."""
        if values.type is bool and values.action != "store_true":
            raise ValueError('If type is bool, action must be "store_true"')
        return values

    @model_validator(mode="after")  # type: ignore
    @classmethod
    def remove_props_on_store_true(cls, values: "ArgparseArgumentModel"):
        """Remove type, nargs, and choices if action is store_true."""
        if values.action == "store_true":
            values.type = None
            values.nargs = None
            values.choices = None
        return values

    # override
    def model_dump(self, **kwargs):
        """Override the model_dump method to remove empty choices."""
        res = super().model_dump(**kwargs)

        # Check if choices is present and if it's an empty tuple remove it
        if "choices" in res and not res["choices"]:
            del res["choices"]

        return res


class ArgparseArgumentGroupModel(BaseModel):
    """Pydantic model for a custom argument group."""

    name: str
    arguments: List[ArgparseArgumentModel]