Add support for overriding param and return types

README.md

@@ -24,7 +24,7 @@ $ ac cli register_app --name weather
 ```
 
 Now, you can call your function from the command-line:
-```
+```bash
 $ ac weather get_weather --location Amsterdam
 21 degrees celsius. Sunny all day in Amsterdam!
 

auto_cli/cli.py

@@ -1,50 +1,50 @@
 import importlib
 import sys
 from pathlib import Path
-from typing import Callable, Dict, List, Optional, TypeVar, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 from .configuration import Configuration
 from .parsing import create_parser
+from .types import Command
 from .utils import _print_and_quit
 
-# In practice, this dict will be overriden by _load_app
-# but it is here to keep the linters happy and for testing
-REGISTERED_COMMANDS: Dict[str, Callable] = {}
-
-ReturnType = TypeVar("ReturnType")
 
-
-def run_func_with_argv(
-    function: Callable[..., ReturnType], argv: List[str], command: str
-) -> ReturnType:
-    parser = create_parser(function, command)
+def run_func_with_argv(command: Command, argv: List[str]) -> Any:
+    parser = create_parser(command)
     args = parser.parse(argv)
-    retval = function(**args)
+    retval = command.function(**args)
+    if command.return_type is not None:
+        retval = command.return_type(retval)
     return retval
 
 
-def run_func_with_argv_and_print(
-    function: Callable[..., ReturnType], argv: List[str], command: str
-) -> None:
-    result = run_func_with_argv(function, argv, command)
+def run_func_with_argv_and_print(command: Command, argv: List[str]) -> None:
+    result = run_func_with_argv(command, argv)
     if result is not None:
         print(result)
 
 
+# In practice, this dict will be overriden by _load_app
+# but it is here to keep the linters happy and for testing
+REGISTERED_COMMANDS: Dict[str, Command] = {}
+
+
 def register_command(
-    function: Union[str, Callable], name: Optional[str] = None
+    function: Union[str, Callable[..., Any]],
+    name: Optional[str] = None,
+    parameter_types: Optional[Dict[str, Callable]] = None,
+    return_type: Optional[Callable[[Any], Any]] = None,
 ) -> None:
     """Register `function` as an available command"""
-    # TODO(joris): Add custom types for arguments
-    # TODO(joris): Add result formatter option
     python_function: Callable
     if isinstance(function, str):
         python_function = _get_function_from_str(function)
     else:
         python_function = function
     command_name = name or python_function.__name__
+    command = Command(command_name, python_function, parameter_types, return_type)
     # REGISTERED_COMMANDS is magically defined by _load_app
-    REGISTERED_COMMANDS[command_name] = python_function
+    REGISTERED_COMMANDS[command_name] = command
 
 
 def register_app(name: str, location: Optional[Path] = None) -> None:
@@ -59,9 +59,9 @@ def run_command(app: str, argv: List[str]) -> None:
     if len(argv) == 0:
         command_help = _command_help(commands)
         _print_and_quit(f"No command given. Available commands:\n{command_help}")
-    command, *argv = argv
-    function = commands[command]
-    run_func_with_argv_and_print(function, argv, command)
+    command_name, *argv = argv
+    command = commands[command_name]
+    run_func_with_argv_and_print(command, argv)
 
 
 def run() -> None:
@@ -88,24 +88,33 @@ def _get_function_from_str(path: str) -> Callable:
     return function
 
 
-def _load_app(name: str) -> Dict[str, Callable]:
+def _load_app(name: str) -> Dict[str, Command]:
     """Load the commands registered in auto_cli.py"""
     with Configuration() as config:
         ac_code = config.get_app_ac_content(name)
-    ac_code = f"""
-import auto_cli
-auto_cli.cli.REGISTERED_COMMANDS = REGISTERED_COMMANDS
-{ac_code}"""
-    registered_commands: Dict[str, Callable] = {}
+
+    # We want ac_code to be executed and fill a variable
+    # that is local to this function with the possible commands.
+    # Hence, we override the global REGISTERED_COMMANDS
+    # with a dictionary that we can return from this function.
+    ac_code = "\n".join(
+        [
+            "import auto_cli",
+            "auto_cli.cli.REGISTERED_COMMANDS = REGISTERED_COMMANDS",
+            ac_code,
+        ]
+    )
+
+    registered_commands: Dict[str, Command] = {}
     exec(ac_code, {"REGISTERED_COMMANDS": registered_commands})
     return registered_commands
 
 
-def _command_help(commands: Dict[str, Callable]) -> str:
+def _command_help(commands: Dict[str, Command]) -> str:
     longest_name = max(map(len, commands))
     return "\n".join(
-        f"{name.ljust(longest_name)}{_function_help(function)}"
-        for name, function in commands.items()
+        f"{name.ljust(longest_name)}{_function_help(command.function)}"
+        for name, command in commands.items()
     )
 
 

auto_cli/parsing.py

@@ -2,6 +2,7 @@
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
+from .types import Command
 from .utils import _print_and_quit
 
 
@@ -27,9 +28,10 @@ def add_param_transformer(self, param_name: str, transform: Callable) -> None:
         self.post_processing[param_name] = transform
 
 
-def create_parser(function: Callable, command_name: str) -> ArgumentParser:
-    """Create a parser for the given function"""
-    parser = ArgumentParser(description=f"{command_name}: {function.__doc__}")
+def create_parser(command: Command) -> ArgumentParser:
+    """Create a parser for the given command"""
+    function = command.function
+    parser = ArgumentParser(description=f"{command.name}: {function.__doc__}")
     signature = inspect.signature(function)
     parameters = dict(signature.parameters)
     argspec = inspect.getfullargspec(function)
@@ -39,12 +41,14 @@ def create_parser(function: Callable, command_name: str) -> ArgumentParser:
     if argspec.varkw is not None:
         del parameters[argspec.varkw]
 
+    param_types = command.parameter_types or {}
     for param_name, parameter in parameters.items():
+        annotation = param_types.get(param_name, parameter.annotation)
         kwargs = {
             "required": not _has_default(parameter),
             "default": parameter.default if _has_default(parameter) else None,
             # The params above might be overwritten by the function below
-            **_get_type_params(parameter.annotation, param_name, function),
+            **_get_type_params(annotation, param_name, function),
         }
 
         parser.add_argument(f"--{param_name}", **kwargs)

auto_cli/tests/test_cli.py

@@ -3,13 +3,14 @@
 import pytest_mock
 
 from auto_cli import cli
+from auto_cli.types import Command
 
 
 def test_run_func_with_argv() -> None:
     def func_to_test(a: int, b: int = 38) -> int:
         return a + b
 
-    result = cli.run_func_with_argv(func_to_test, ["--a", "4"], func_to_test.__name__)
+    result = cli.run_func_with_argv(Command.from_func(func_to_test), ["--a", "4"])
     assert result == 42
 
 
@@ -22,11 +23,9 @@ def test_register_command() -> None:
 
 def test_run_command(mocker: pytest_mock.MockFixture) -> None:
     mocked_cmd = mocker.MagicMock()
+    registered_commands = {"test_cmd": Command.from_func(mocked_cmd, name="test_cmd")}
+    mocker.patch("auto_cli.cli._load_app", return_value=registered_commands)
 
-    def register_func(_: str):
-        cli.REGISTERED_COMMANDS["test_cmd"] = mocked_cmd
-
-    mocker.patch("auto_cli.cli._load_app", return_value={"test_cmd": mocked_cmd})
     cli.run_command("my_app", ["test_cmd"])
     mocked_cmd.assert_called_once()
 

auto_cli/tests/test_parsing.py

@@ -1,13 +1,14 @@
 from typing import List, Tuple
 
 from auto_cli.parsing import create_parser
+from auto_cli.types import Command
 
 
 def test_create_parser_simple() -> None:
     def func_to_test(a: int, b: str) -> int:
         return a + len(b)
 
-    parser = create_parser(func_to_test, "")
+    parser = create_parser(Command.from_func(func_to_test))
     args = parser.parse(["--a", "42", "--b", "1234"])
     assert args == {"a": 42, "b": "1234"}
 
@@ -16,7 +17,7 @@ def test_create_parser_defaults() -> None:
     def func_to_test(a: int, b: int = 38) -> int:
         return a + b
 
-    parser = create_parser(func_to_test, "")
+    parser = create_parser(Command.from_func(func_to_test))
     args_no_default = parser.parse(["--a", "1", "--b", "42"])
     assert args_no_default == {"a": 1, "b": 42}
 
@@ -28,7 +29,7 @@ def test_create_parser_bool() -> None:
     def func_to_test(a: bool) -> bool:
         return a
 
-    parser = create_parser(func_to_test, "")
+    parser = create_parser(Command.from_func(func_to_test))
     args_with_flag = parser.parse(["--a"])
     assert args_with_flag == {"a": True}
 
@@ -40,7 +41,7 @@ def test_create_parser_list() -> None:
     def func_to_test(a: List[int]) -> int:
         return sum(a)
 
-    parser = create_parser(func_to_test, "")
+    parser = create_parser(Command.from_func(func_to_test))
     nums = [1, 3, 5, 7]
     args = parser.parse(["--a"] + list(map(str, nums)))
 
@@ -51,8 +52,18 @@ def test_create_parser_tuple() -> None:
     def func_to_test(a: Tuple[int, int], b: bool) -> int:
         return sum(a)
 
-    parser = create_parser(func_to_test, "")
+    parser = create_parser(Command.from_func(func_to_test))
     nums = (42, 1337)
     args = parser.parse(["--a"] + list(map(str, nums)))
 
     assert args == {"a": nums, "b": False}
+
+
+def test_override_param_type() -> None:
+    def func_to_test(a: int, b):
+        return a + b
+
+    command = Command("cmd_name", func_to_test, {"b": int}, None)
+    parser = create_parser(command)
+    args = parser.parse(["--a", "4", "--b", "5"])
+    assert args == {"a": 4, "b": 5}

auto_cli/types.py

@@ -0,0 +1,15 @@
+from typing import Any, Callable, Dict, NamedTuple, Optional
+
+
+class Command(NamedTuple):
+    name: str
+    function: Callable[..., Any]
+    parameter_types: Optional[Dict[str, Callable]]
+    return_type: Optional[Callable[[Any], Any]]
+
+    @staticmethod
+    def from_func(function: Callable, name: Optional[str] = None) -> "Command":
+        """Convience function for what should be the most common case:
+        Creating a Command from a Python function with no bells or whistles
+        """
+        return Command(name or function.__name__, function, None, None)