Skip to content

Commit

Permalink
[Feature] Handle repeated non standard arguments (#6366)
Browse files Browse the repository at this point in the history
* remove about message

* handle repeated arguments

* accomodate all the choices from different providers

---------

Co-authored-by: Igor Radovanovic <74266147+IgorWounds@users.noreply.github.com>
  • Loading branch information
hjoaquim and IgorWounds authored May 7, 2024
1 parent 5c0b36c commit 756eebd
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 20 deletions.
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

0 comments on commit 756eebd

Please sign in to comment.