From fcf519da7f3642c2c8f4e2d518188d43a5e459df Mon Sep 17 00:00:00 2001 From: hjoaquim Date: Wed, 9 Aug 2023 14:14:59 +0100 Subject: [PATCH] sdk powered terminal commands with clean help messages --- argparse_translator/argparse_translator.py | 51 +++++++++++++++++----- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/argparse_translator/argparse_translator.py b/argparse_translator/argparse_translator.py index d39f36b4fd0..9a18d24ad46 100644 --- a/argparse_translator/argparse_translator.py +++ b/argparse_translator/argparse_translator.py @@ -20,14 +20,11 @@ from typing import ( from pydantic import BaseModel from typing_extensions import Annotated +from openbb_core.app.model.custom_parameter import OpenBBCustomParameter + SEP = "__" -class OpenBBCustomArgument(BaseModel): - help: Optional[str] = None - # action: Optional[Any] = None - - class ArgparseActionType(Enum): store = "store" store_true = "store_true" @@ -48,7 +45,7 @@ class ArgparseTranslator: self._parser = argparse.ArgumentParser( prog=func.__name__, - description=func.__doc__, + description=self._build_description(func.__doc__), formatter_class=argparse.RawTextHelpFormatter, add_help=add_help, ) @@ -59,11 +56,26 @@ class ArgparseTranslator: 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 ----------"] + + 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: @@ -73,6 +85,7 @@ class ArgparseTranslator: 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) @@ -105,27 +118,32 @@ class ArgparseTranslator: @staticmethod def _split_annotation( base_annotation: Type[Any], - ) -> tuple[Type[Any], List[OpenBBCustomArgument]]: + ) -> Tuple[Type[Any], List[OpenBBCustomParameter]]: + """Find the base annotation and the custom annotations, namely the OpenBBCustomParameter.""" 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, OpenBBCustomArgument) + if isinstance(annotation, OpenBBCustomParameter) ] @classmethod def _get_argument_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) - help_annotation = custom_annotations[0].help if custom_annotations else None + 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 def _get_nargs(self, param: inspect.Parameter) -> Optional[str]: + """Returns the nargs annotation for the given parameter.""" param_type = self.type_hints[param.name] if get_origin(param_type) is list: @@ -133,6 +151,7 @@ class ArgparseTranslator: 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: @@ -162,7 +181,7 @@ class ArgparseTranslator: name=f"{param.name}{SEP}{child_param.name}", annotation=Annotated[ child_param.annotation, - OpenBBCustomArgument( + OpenBBCustomParameter( help=param_type.schema()["properties"][ child_param.name ].get("help", None) @@ -207,6 +226,7 @@ class ArgparseTranslator: @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: @@ -222,6 +242,7 @@ class ArgparseTranslator: 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(): @@ -239,6 +260,16 @@ class ArgparseTranslator: 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)