Python函数作为命令行使用

func_cmdline.py

#!/usr/bin/env python3
# coding: utf-8

__author__ = "ChenyangGao <https://chenyanggao.github.io/>"
__version__ = (0, 1, 1)
__all__ = ["func_cmdline"]

# Inspired by:
# - https://pypi.org/project/typer/

import inspect

from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
from functools import partial
from textwrap import dedent, indent
from typing import cast, Callable, Optional


def func_cmdline(
    fn: Optional[Callable] = None, 
    /, 
    parser: Optional[ArgumentParser] = None, 
    parse_map: dict[type, Callable] = {}, 
    helpinfo_map: dict[str, str] = {}, 
) -> Callable:
    """装饰器让一个函数，可以作为命令行使用

    :param fn: 被装饰的函数
    :param parser: 命令行解析器对象
    :param parse_map: 不同类型的解析函数的映射，
        字符串通过相应的解析函数来转换成对应对象，
        如果没有相应的解析函数，则尝试用对象本身
        （它的构造器）
    :param helpinfo_map: 不同参数的帮助说明，
        用于命令行的帮助信息
    :param: 返函数`fn`本身，只是它身上绑定了几个属性
        .parser 字段，使用的命令行解析器对象
        .parse() 方法，解析命令行参数
        .run() 方法，解析命令行参数并运行函数`fn`，
            返回函数`fn`的执行结果

    :return: 被装饰函数本身，添加了几个属性
    """
    if fn is None:
        return partial(
            func_cmdline, 
            parser=parser, 
            parse_map=parse_map, 
            helpinfo_map=helpinfo_map, 
        )
    fn = cast(Callable, fn)
    params = inspect.signature(fn).parameters
    if parser is None:
        parser = ArgumentParser(
            formatter_class=RawDescriptionHelpFormatter)
    parser = cast(ArgumentParser, parser)
    if parser.epilog is None:
        parser.epilog = ""
        if fn.__doc__ is not None:
            parser.epilog += "documentation:\n" + \
                indent(fn.__doc__, "    ")
        parser.epilog += "\n\nsource code:\n" + \
            indent(dedent(inspect.getsource(fn)), "    ")
    abbrs: dict[str, int] = {}
    for name, param in params.items():
        kargs: dict = {"dest": name, "help": helpinfo_map.get(name)}
        kind = param.kind
        if kind is inspect._ParameterKind.VAR_POSITIONAL:
            kargs["required"] = False
            kargs["nargs"] = "*"
            kargs["default"] = []
        elif kind is inspect._ParameterKind.VAR_KEYWORD:
            continue
        default = param.default
        if default is inspect._empty:
            kargs.setdefault("required", True)
        else:
            kargs["default"] = default
        anno = param.annotation
        if anno is not inspect._empty:
            if anno is bool:
                kargs["required"] = False
                if "default" in kargs:
                    if kargs["default"]:
                        kargs["action"] = "store_false"
                    else:
                        kargs["action"] = "store_true"
                else:
                    kargs["action"] = "store_true"
            elif anno is list:
                kargs["nargs"] = "*"
            elif getattr(anno, "__origin__", None) is list and len(anno.__args__) == 1:
                kargs["nargs"] = "*"
                kargs["type"] = parse_map.get(anno.__args__[0], anno.__args__[0])
            else:
                kargs["type"] = parse_map.get(anno, anno)
        names = []
        first_letter = name[0]
        n = abbrs.setdefault(first_letter, 0)
        if n == 0:
            abbr = first_letter
        else:
            abbr = f"{first_letter}{n}"
            abbrs[first_letter] += 1
        names.append(f"-{abbr}")
        names.append(f"--{name}")
        if "_" in name:
            names.append(f"--{name.replace('_', '-')}")
        parser.add_argument(*names, **kargs)
    def parse(argv: Optional[list[str]] = None) -> Namespace:
        nonlocal parser
        parser = cast(ArgumentParser, parser)
        return parser.parse_args(argv)
    def run(argv: Optional[list[str]] = None):
        "解析命令行参数，默认使用"
        nonlocal fn, parser
        fn = cast(Callable, fn)
        parser = cast(ArgumentParser, parser)
        args = parser.parse_args(argv).__dict__
        pargs = []
        kargs = {}
        for name, param in params.items():
            kind = param.kind
            if name in args:
                if kind.value < 2:
                    pargs.append(args[name])
                elif kind.value == 2:
                    pargs.extend(args[name])
                else:
                    kargs[name] = args[name]
        return fn(*pargs, **kargs)
    setattr(fn, "parser", parser)
    setattr(fn, "parse", parse)
    setattr(fn, "run", run)
    return fn


if __name__ == "__main__":
    from datetime import datetime

    @func_cmdline(
        parse_map={
            datetime: lambda s: datetime.strptime(s, "%Y-%m-%d %H:%M:%S"), 
        }, 
        helpinfo_map=dict(
            a="参数a", b="参数b", c="参数c", 
            d="参数d", e="参数e", 
        )
    )
    def foo(
        a: int, /, b: bool = False, *c: float, 
        d: list[datetime], e: complex = 4, 
    ):
        "Example for function as command line."
        print(locals())

    foo.run()