Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Handle repeated non standard arguments #6366

Merged
merged 9 commits into from
May 7, 2024
105 changes: 90 additions & 15 deletions cli/openbb_cli/argparse_translator/argparse_translator.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import argparse
import inspect
import re
from copy import deepcopy
from enum import Enum
from typing import (
Expand All @@ -21,6 +22,8 @@
from pydantic import BaseModel, model_validator
from typing_extensions import Annotated

# pylint: disable=protected-access

SEP = "__"


Expand Down Expand Up @@ -200,7 +203,7 @@ def __init__(
self.func = func
self.signature = inspect.signature(func)
self.type_hints = get_type_hints(func)
self.provider_parameters = []
self.provider_parameters: List[str] = []

self._parser = argparse.ArgumentParser(
prog=func.__name__,
Expand All @@ -217,23 +220,95 @@ def __init__(
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"}, exclude_none=True)

# If the argument is already in use, we can't repeat it
if f"--{argument.name}" not in self._parser_arguments():
argparse_group.add_argument(f"--{argument.name}", **kwargs)
self.provider_parameters.append(argument.name)
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_optional_arguments(arg):
for action_group in self._parser._action_groups:
if action_group.title == "optional arguments":
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

def _parser_arguments(self) -> List[str]:
"""Get all the arguments from all groups currently defined on the parser."""
arguments_in_use: List[str] = []
# 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)
self.provider_parameters.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 optional arguments
if _in_optional_arguments(argument.name):
for action in self._parser._actions:
if action.dest == argument.name:
# update choices
action.choices = choices
# update help
action.help = _update_providers(
action.help or "", [group.title]
)
return

# pylint: disable=protected-access
for action_group in self._parser._action_groups:
for action in action_group._group_actions:
arguments_in_use.extend(action.option_strings)
# 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

return arguments_in_use
# 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:
Expand Down
5 changes: 0 additions & 5 deletions cli/openbb_cli/controllers/base_controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -908,11 +908,6 @@ def parse_known_args_and_warn(

if "--help" in other_args or "-h" in other_args:
txt_help = parser.format_help() + "\n"
if parser.prog != "about":
txt_help += (
f"For more information and examples, use 'about {parser.prog}' "
f"to access the related guide.\n"
)
session.console.print(f"[help]{txt_help}[/help]")
return None

Expand Down
Loading