openbb-cli 1.0.0__py3-none-any.whl

openbb_cli/argparse_translator/argparse_translator.py

@@ -0,0 +1,627 @@
+import argparse
+import inspect
+import re
+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
+
+# pylint: disable=protected-access
+
+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[Tuple]
+
+    @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, custom_choices: Any) -> 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_)
+
+        if type_origin is list:
+            type_ = get_args(type_)[0]
+
+            if get_origin(type_) is Literal:
+                choices = get_args(type_)
+
+        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_)
+
+        if custom_choices:
+            return tuple(custom_choices)
+
+        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"], custom_choices=arg["choices"]
+                            ),
+                        )
+                    )
+
+                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: Dict[str, List[str]] = {}
+
+        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:
+                self.provider_parameters[group.name] = []
+                argparse_group = self._parser.add_argument_group(group.name)
+                for argument in group.arguments:
+                    self._handle_argument_in_groups(argument, argparse_group)
+
+    def _handle_argument_in_groups(self, argument, group):
+        """Handle the argument and add it to the parser."""
+
+        def _in_group(arg, group_title):
+            for action_group in self._parser._action_groups:
+                if action_group.title == group_title:
+                    for action in action_group._group_actions:
+                        opts = action.option_strings
+                        if (opts and opts[0] == arg) or action.dest == arg:
+                            return True
+            return False
+
+        def _remove_argument(arg) -> List[Optional[str]]:
+            groups_w_arg = []
+
+            # remove the argument from the parser
+            for action in self._parser._actions:
+                opts = action.option_strings
+                if (opts and opts[0] == arg) or action.dest == arg:
+                    self._parser._remove_action(action)
+                    break
+
+            # remove from all groups
+            for action_group in self._parser._action_groups:
+                for action in action_group._group_actions:
+                    opts = action.option_strings
+                    if (opts and opts[0] == arg) or action.dest == arg:
+                        action_group._group_actions.remove(action)
+                        groups_w_arg.append(action_group.title)
+
+            # remove from _action_groups dict
+            self._parser._option_string_actions.pop(f"--{arg}", None)
+
+            return groups_w_arg
+
+        def _get_arg_choices(arg) -> Tuple:
+            for action in self._parser._actions:
+                opts = action.option_strings
+                if (opts and opts[0] == arg) or action.dest == arg:
+                    return tuple(action.choices or ())
+            return ()
+
+        def _update_providers(
+            input_string: str, new_provider: List[Optional[str]]
+        ) -> str:
+            pattern = r"\(provider:\s*(.*?)\)"
+            providers = re.findall(pattern, input_string)
+            providers.extend(new_provider)
+            # remove pattern from help and add with new providers
+            input_string = re.sub(pattern, "", input_string).strip()
+            return f"{input_string} (provider: {', '.join(providers)})"
+
+        # check if the argument is already in use, if not, add it
+        if f"--{argument.name}" not in self._parser._option_string_actions:
+            kwargs = argument.model_dump(exclude={"name"}, exclude_none=True)
+            group.add_argument(f"--{argument.name}", **kwargs)
+            if group.title in self.provider_parameters:
+                self.provider_parameters[group.title].append(argument.name)
+
+        else:
+            kwargs = argument.model_dump(exclude={"name"}, exclude_none=True)
+            model_choices = kwargs.get("choices", ()) or ()
+            # extend choices
+            choices = tuple(set(_get_arg_choices(argument.name) + model_choices))
+
+            # check if the argument is in the required arguments
+            if _in_group(argument.name, group_title="required arguments"):
+                for action in self._required._group_actions:
+                    if action.dest == argument.name and choices:
+                        # update choices
+                        action.choices = choices
+                return
+
+            # check if the argument is in the optional arguments
+            if _in_group(argument.name, group_title="optional arguments"):
+                for action in self._parser._actions:
+                    if action.dest == argument.name:
+                        # update choices
+                        if choices:
+                            action.choices = choices
+                        if argument.name not in self.signature.parameters:
+                            # update help
+                            action.help = _update_providers(
+                                action.help or "", [group.title]
+                            )
+                return
+
+            # if the argument is in use, remove it from all groups
+            # and return the groups that had the argument
+            groups_w_arg = _remove_argument(argument.name)
+            groups_w_arg.append(group.title)  # add current group
+
+            # add it to the optional arguments group instead
+            if choices:
+                kwargs["choices"] = choices  # update choices
+            # add provider info to the help
+            kwargs["help"] = _update_providers(argument.help or "", groups_w_arg)
+            self._parser.add_argument(f"--{argument.name}", **kwargs)
+
+    @property
+    def parser(self) -> argparse.ArgumentParser:
+        return deepcopy(self._parser)
+
+    @staticmethod
+    def _build_description(func_doc: str) -> str:
+        """Builds the description of the argparse program from the function docstring."""
+
+        patterns = ["openbb\n        ======", "Parameters\n        ----------"]
+
+        if func_doc:
+            for pattern in patterns:
+                if pattern in func_doc:
+                    func_doc = func_doc[: func_doc.index(pattern)].strip()
+                    break
+
+        return func_doc
+
+    @staticmethod
+    def _param_is_default(param: inspect.Parameter) -> bool:
+        """Returns True if the parameter has a default value."""
+        return param.default != inspect.Parameter.empty
+
+    def _get_action_type(self, param: inspect.Parameter) -> str:
+        """Returns the argparse action type for the given parameter."""
+        param_type = self.type_hints[param.name]
+
+        if param_type == bool:
+            return ArgparseActionType.store_true.value
+        return ArgparseActionType.store.value
+
+    def _get_type_and_choices(
+        self, param: inspect.Parameter
+    ) -> Tuple[Type[Any], Tuple[Any, ...]]:
+        """Returns the type and choices for the given parameter."""
+        param_type = self.type_hints[param.name]
+        type_origin = get_origin(param_type)
+
+        choices = ()
+
+        if type_origin is Literal:
+            choices = get_args(param_type)
+            param_type = type(choices[0])  # type: ignore
+
+        if type_origin is list:  # TODO: dict should also go here
+            param_type = get_args(param_type)[0]
+
+            if get_origin(param_type) is Literal:
+                choices = get_args(param_type)
+                param_type = type(choices[0])  # type: ignore
+
+        if type_origin is Union:
+            union_args = get_args(param_type)
+            if str in union_args:
+                param_type = str
+
+            # check if it's an Optional, which would be a Union with NoneType
+            if type(None) in get_args(param_type):
+                # remove NoneType from the args
+                args = [arg for arg in get_args(param_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"
+                    )
+                param_type = args[0]
+
+                if get_origin(param_type) is Literal:
+                    choices = get_args(param_type)
+                    param_type = type(choices[0])  # type: ignore
+
+        # if there are custom choices, override
+        choices = self._get_argument_custom_choices(param) or choices  # type: ignore
+
+        return param_type, choices
+
+    @staticmethod
+    def _split_annotation(
+        base_annotation: Type[Any], custom_annotation_type: Type
+    ) -> Tuple[Type[Any], List[Any]]:
+        """Find the base annotation and the custom annotations, namely the OpenBBField."""
+        if get_origin(base_annotation) is not Annotated:
+            return base_annotation, []
+        base_annotation, *maybe_custom_annotations = get_args(base_annotation)
+        return base_annotation, [
+            annotation
+            for annotation in maybe_custom_annotations
+            if isinstance(annotation, custom_annotation_type)
+        ]
+
+    @classmethod
+    def _get_argument_custom_help(cls, param: inspect.Parameter) -> Optional[str]:
+        """Returns the help annotation for the given parameter."""
+        base_annotation = param.annotation
+        _, custom_annotations = cls._split_annotation(base_annotation, OpenBBField)
+        help_annotation = (
+            custom_annotations[0].description if custom_annotations else None
+        )
+        if not help_annotation:
+            # try to get it from the docstring
+            pass
+        return help_annotation
+
+    @classmethod
+    def _get_argument_custom_choices(cls, param: inspect.Parameter) -> Optional[str]:
+        """Returns the help annotation for the given parameter."""
+        base_annotation = param.annotation
+        _, custom_annotations = cls._split_annotation(base_annotation, OpenBBField)
+        choices_annotation = (
+            custom_annotations[0].choices if custom_annotations else None
+        )
+        if not choices_annotation:
+            # try to get it from the docstring
+            pass
+        return choices_annotation
+
+    def _get_nargs(self, param: inspect.Parameter) -> Optional[str]:
+        """Returns the nargs annotation for the given parameter."""
+        param_type = self.type_hints[param.name]
+        origin = get_origin(param_type)
+
+        if origin is list:
+            return "+"
+
+        if origin is Union and any(
+            get_origin(arg) is list for arg in get_args(param_type)
+        ):
+            return "+"
+
+        return None
+
+    def _generate_argparse_arguments(self, parameters) -> None:
+        """Generates the argparse arguments from the function parameters."""
+        for param in parameters.values():
+            # TODO : how to handle kwargs?
+            # it's possible to add unknown arguments when parsing as follows:
+            # args, unknown_args = parser.parse_known_args()
+            if param.name == "kwargs":
+                continue
+
+            param_type, choices = self._get_type_and_choices(param)
+
+            # if the param is a custom type, we need to flatten it
+            if inspect.isclass(param_type) and issubclass(param_type, BaseModel):
+                # update type hints with the custom type fields
+                type_hints = get_type_hints(param_type)
+                # prefix the type hints keys with the param name
+                type_hints = {
+                    f"{param.name}{SEP}{key}": value
+                    for key, value in type_hints.items()
+                }
+                self.type_hints.update(type_hints)
+                # create a signature from the custom type
+                sig = inspect.signature(param_type)
+
+                # add help to the annotation
+                annotated_parameters: List[inspect.Parameter] = []
+                for child_param in sig.parameters.values():
+                    new_child_param = child_param.replace(
+                        name=f"{param.name}{SEP}{child_param.name}",
+                        annotation=Annotated[
+                            child_param.annotation,
+                            OpenBBField(
+                                description=param_type.model_json_schema()[
+                                    "properties"
+                                ][child_param.name].get("description", None)
+                            ),
+                        ],
+                        kind=inspect.Parameter.KEYWORD_ONLY,
+                    )
+                    annotated_parameters.append(new_child_param)
+
+                # replacing with the annotated parameters
+                new_signature = inspect.Signature(
+                    parameters=annotated_parameters,
+                    return_annotation=sig.return_annotation,
+                )
+                self._generate_argparse_arguments(new_signature.parameters)
+
+                # the custom type itself should not be added as an argument
+                continue
+
+            required = not self._param_is_default(param)
+
+            kwargs = {
+                "type": param_type,
+                "dest": param.name,
+                "default": param.default,
+                "required": required,
+                "action": self._get_action_type(param),
+                "help": self._get_argument_custom_help(param),
+                "nargs": self._get_nargs(param),
+            }
+
+            if choices:
+                kwargs["choices"] = choices
+
+            if param_type == bool:
+                # store_true action does not accept the below kwargs
+                kwargs.pop("type")
+                kwargs.pop("nargs")
+
+            if required:
+                self._required.add_argument(
+                    f"--{param.name}",
+                    **kwargs,
+                )
+            else:
+                self._parser.add_argument(
+                    f"--{param.name}",
+                    **kwargs,
+                )
+
+    @staticmethod
+    def _unflatten_args(args: dict) -> Dict[str, Any]:
+        """Unflatten the args that were flattened by the custom types."""
+        result: Dict[str, Any] = {}
+        for key, value in args.items():
+            if SEP in key:
+                parts = key.split(SEP)
+                nested_dict = result
+                for part in parts[:-1]:
+                    if part not in nested_dict:
+                        nested_dict[part] = {}
+                    nested_dict = nested_dict[part]
+                nested_dict[parts[-1]] = value
+            else:
+                result[key] = value
+        return result
+
+    def _update_with_custom_types(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Update the kwargs with the custom types."""
+        # for each argument in the signature that is a custom type, we need to
+        # update the kwargs with the custom type kwargs
+        for param in self.signature.parameters.values():
+            # TODO : how to handle kwargs?
+            if param.name == "kwargs":
+                continue
+            param_type, _ = self._get_type_and_choices(param)
+            if inspect.isclass(param_type) and issubclass(param_type, BaseModel):
+                custom_type_kwargs = kwargs[param.name]
+                kwargs[param.name] = param_type(**custom_type_kwargs)
+
+        return kwargs
+
+    def execute_func(
+        self,
+        parsed_args: Optional[argparse.Namespace] = None,
+    ) -> Any:
+        """
+        Executes the original function with the parsed arguments.
+
+        Args:
+            parsed_args (Optional[argparse.Namespace], optional): The parsed arguments. Defaults to None.
+
+        Returns:
+            Any: The return value of the original function.
+
+        """
+        kwargs = self._unflatten_args(vars(parsed_args))
+        kwargs = self._update_with_custom_types(kwargs)
+
+        provider = kwargs.get("provider")
+        provider_args = []
+        if provider and provider in self.provider_parameters:
+            provider_args = self.provider_parameters[provider]
+        else:
+            for args in self.provider_parameters.values():
+                provider_args.extend(args)
+
+        # remove kwargs that doesn't match the signature or provider parameters
+        kwargs = {
+            key: value
+            for key, value in kwargs.items()
+            if key in self.signature.parameters or key in provider_args
+        }
+
+        return self.func(**kwargs)
+
+    def parse_args_and_execute(self) -> Any:
+        """
+        Parses the arguments and executes the original function.
+
+        Returns:
+            Any: The return value of the original function.
+        """
+        parsed_args = self._parser.parse_args()
+        return self.execute_func(parsed_args)
+
+    def translate(self) -> Callable:
+        """
+        Wraps the original function with an argparse program.
+
+        Returns:
+            Callable: The original function wrapped with an argparse program.
+        """
+
+        def wrapper_func():
+            return self.parse_args_and_execute()
+
+        return wrapper_func