From 6d2dc0aaab8ab5249e3bb91a736d489d8e17a5f8 Mon Sep 17 00:00:00 2001 From: hjoaquim Date: Tue, 20 Jun 2023 09:55:58 +0100 Subject: [PATCH] branch with the argparse translator to be used with the sdk core --- argparse_translator/__init__.py | 0 argparse_translator/argparse_translator.py | 194 +++++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 argparse_translator/__init__.py create mode 100644 argparse_translator/argparse_translator.py diff --git a/argparse_translator/__init__.py b/argparse_translator/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/argparse_translator/argparse_translator.py b/argparse_translator/argparse_translator.py new file mode 100644 index 00000000000..65b0657eb1a --- /dev/null +++ b/argparse_translator/argparse_translator.py @@ -0,0 +1,194 @@ +import argparse +import inspect +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Type, + get_args, + get_origin, + get_type_hints, +) + +from pydantic import BaseModel +from typing_extensions import Annotated + +SEP = "__" + + +class ArgparseCustomArgument(BaseModel): + help: Optional[str] = None + # action: Optional[Any] = None + + +class ArgparseActionType(Enum): + store = "store" + store_true = "store_true" + + +class ArgparseTranslator: + def __init__(self, func: Callable): + """ + Initializes the ArgparseTranslator. + + Args: + func (Callable): The function to translate into an argparse program. + """ + self.func = func + self.signature = inspect.signature(func) + self.type_hints = get_type_hints(func) + + self.parser = argparse.ArgumentParser( + description=func.__doc__, formatter_class=argparse.RawTextHelpFormatter + ) + + @staticmethod + def _param_is_default(param: inspect.Parameter) -> bool: + return param.default != inspect.Parameter.empty + + def _get_action_type(self, param: inspect.Parameter) -> str: + param_type = self.type_hints[param.name] + + if param_type == bool: + return ArgparseActionType.store_true.value + return ArgparseActionType.store.value + + def _get_type(self, param: inspect.Parameter): + param_type = self.type_hints[param.name] + + if get_origin(param_type) is list: # TODO: dict should also go here + return get_args(param_type)[0] + + return param_type + + @staticmethod + def _split_annotation( + base_annotation: Type[Any], + ) -> tuple[Type[Any], List[ArgparseCustomArgument]]: + 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, ArgparseCustomArgument) + ] + + @classmethod + def _get_argument_help(cls, param: inspect.Parameter) -> Optional[str]: + base_annotation = param.annotation + _, custom_annotations = cls._split_annotation(base_annotation) + help_annotation = custom_annotations[0].help 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]: + param_type = self.type_hints[param.name] + + if get_origin(param_type) is list: + return "+" + return None + + def _generate_argparse_arguments(self, parameters) -> None: + for param in parameters.values(): + param_type = self._get_type(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(): + child_param = child_param.replace( + name=f"{param.name}{SEP}{child_param.name}", + annotation=Annotated[ + child_param.annotation, + ArgparseCustomArgument( + help=param_type.schema()["properties"][ + child_param.name + ].get("help", None) + ), + ], + kind=inspect.Parameter.KEYWORD_ONLY, + ) + annotated_parameters.append(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 + + self.parser.add_argument( + f"--{param.name}", + type=param_type, + dest=param.name, + default=param.default, + required=not self._param_is_default(param), + action=self._get_action_type(param), + help=self._get_argument_help(param), + nargs=self._get_nargs(param), + ) + + @staticmethod + def _unflatten_args(args: dict) -> Dict[str, Any]: + 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]: + # 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(): + param_type = self._get_type(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 translate(self) -> Callable: + """ + Wraps the original function with an argparse program. + + Returns: + Callable: The original function wrapped with an argparse program. + """ + self._generate_argparse_arguments(self.signature.parameters) + + def wrapper_func(): + args = self.parser.parse_args() + kwargs = self._unflatten_args(vars(args)) + kwargs = self._update_with_custom_types(kwargs) + return self.func(**kwargs) + + return wrapper_func