path: root/cli/openbb_cli/argparse_translator/argparse_translator.py

import argparse
import inspect
from copy import deepcopy
from enum import Enum
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    Type,
    Union,
    get_args,
    get_origin,
    get_type_hints,
)

from openbb_core.app.model.field import OpenBBField
from pydantic import BaseModel, model_validator
from typing_extensions import Annotated

SEP = "__"


class ArgparseActionType(Enum):
    store = "store"
    store_true = "store_true"


class CustomArgument(BaseModel):

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

    @model_validator(mode="after")  # type: ignore
    @classmethod
    def validate_action(cls, values: "CustomArgument"):
        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: "CustomArgument"):
        if values.action == "store_true":
            values.type = None
            values.nargs = None
            values.choices = None
        return values

    # override
    def model_dump(self, **kwargs):

        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 CustomArgumentGroup(BaseModel):
    name: str
    arguments: List[CustomArgument]


class ReferenceToCustomArgumentsProcessor:
    def __init__(self, reference: Dict[str, Dict]):
        """Initializes the ReferenceToCustomArgumentsProcessor."""
        self.reference = reference
        self.custom_groups: Dict[str, List[CustomArgumentGroup]] = {}

        self.build_custom_groups()

    @staticmethod
    def _make_type_parsable(type_: str) -> type:
        """Make the type parsable by removing the annotations."""
        if "Union" in type_ and "str" in type_:
            return str
        if "Union" in type_ and "int" in type_:
            return int
        if type_ in ["date", "datetime.time", "time"]:
            return str

        if any(x in type_ for x in ["gt=", "ge=", "lt=", "le="]):
            if "Annotated" in type_:
                type_ = type_.replace("Annotated[", "").replace("]", "")
            type_ = type_.split(",")[0]

        return eval(type_)  # noqa: S307, E501 pylint: disable=eval-used

    def _parse_type(self, type_: str) -> type:
        """Parse the type from the string representation."""
        type_ = self._make_type_parsable(type_)  # type: ignore

        if get_origin(type_) is Literal:
            type_ = type(get_args(type_)[0])  # type: ignore

        return type_  # type: ignore

    def _get_nargs(self, type_: type) -> Optional[Union[int, str]]:
        """Get the nargs for the given type."""
        if get_origin(type_) is list:
            return "+"
        return None

    def _get_choices(self, type_: str) -> Tuple:
        """Get the choices for the given type."""
        type_ = self._make_type_parsable(type_)  # type: ignore
        type_origin = get_origin(type_)

        choices = ()

        if type_origin is Literal:
            choices = get_args(type_)
            # param_type = type(choices[0])

        if type_origin is list:
            type_ = get_args(type_)[0]

            if get_origin(type_) is Literal:
                choices = get_args(type_)
                # param_type = type(choices[0])

        if type_origin is Union and type(None) in get_args(type_):
            # remove NoneType from the args
            args = [arg for arg in get_args(type_) if arg != type(None)]
            # if there is only one arg left, use it
            if len(args) > 1:
                raise ValueError("Union with NoneType should have only one type left")
            type_ = args[0]

            if get_origin(type_) is Literal:
                choices = get_args(type_)
                # param_type = type(choices[0])

        return choices

    def build_custom_groups(self):
        """Build the custom groups from the reference."""
        for route, v in self.reference.items():

            for provider, args in v["parameters"].items():
                if provider == "standard":
                    continue

                custom_arguments = []
                for arg in args:
                    if arg.get("standard"):
                        continue

                    type_ = self._parse_type(arg["type"])

                    custom_arguments.append(
                        CustomArgument(
                            name=arg["name"],
                            type=type_,
                            dest=arg["name"],
                            default=arg["default"],
                            required=not (arg["optional"]),
                            action="store" if type_ != bool else "store_true",
                            help=arg["description"],
                            nargs=self._get_nargs(type_),  # type: ignore
                            choices=self._get_choices(arg["type"]),
                        )
                    )

                group = CustomArgumentGroup(name=provider, arguments=custom_arguments)

                if route not in self.custom_groups:
                    self.custom_groups[route] = []

                self.custom_groups[route].append(group)


class ArgparseTranslator:
    def __init__(
        self,
        func: Callable,
        custom_argument_groups: Optional[List[CustomArgumentGroup]] = None,
        add_help: Optional[bool] = True,
    ):
        """
        Initializes the ArgparseTranslator.

        Args:
            func (Callable): The function to translate into an argparse program.
            add_help (Optional[bool], optional): Whether to add the help argument. Defaults to False.
        """
        self.func = func
        self.signature = inspect.signature(func)
        self.type_hints = get_type_hints(func)
        self.provider_parameters = []

        self._parser = argparse.ArgumentParser(
            prog=func.__name__,
            description=self._build_description(func.__doc__),  # type: ignore
            formatter_class=argparse.RawTextHelpFormatter,
            add_help=add_help if add_help else False,
        )
        self._required = self._parser.add_argument_group("required arguments")

        if any(param in self.type_hints for param in self.signature.parameters):
            self._generate_argparse_arguments(self.signature.parameters)

        if custom_argument_groups:
            for group in custom_argument_groups:
                argparse_group = self._parser.add_argument_group(group.name)
                for argument in group.arguments:
                    kwargs = argument.model_dump(exclude={"name"}, e