Please enable Javascript to view the contents

Python函数当命令行使用

 ·  ☕ 4 分钟  ·  🐵 chenyanggao

Python函数当命令行使用

背景

我最近了解了一个快速构建命令行的工具 typer,受到了一些启发。
它可以把一组带有类型注解的 Python 函数,快速地转换成命令行工具。我尝试只用 Python 的内建包,不引入第三方依赖,也实现一个类似的命令行快速构建工具。代码会持续更新,之后会加入更强的功能。

代码实现

注意:使用下面的代码,请用 Python 3.8 或者以上版本。

文件名称是 func_cmdline.pyPython 实现代码如下:

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
#!/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()

使用说明

模块 func_cmdline.py 自带了一个实例:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
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 = True, *c: float, 
    d: list[datetime], e: complex = 1j, 
):
    "Example for function as command line."
    print(locals())

foo.run()

直接运行模块,输出如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
$ python func_cmdline.py -h
usage: func_cmdline.py [-h] -a A [-b] [-c [C ...]] -d [D ...] [-e E]

optional arguments:
  -h, --help            show this help message and exit
  -a A, --a A           参数a
  -b, --b               参数b
  -c [C ...], --c [C ...]
                        参数c
  -d [D ...], --d [D ...]
                        参数d
  -e E, --e E           参数e

documentation:
    Example for function as command line.

source code:
    @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 = 1j,
    ):
        "Example for function as command line."
        print(locals())

随便传入一些必需的参数,不传可选参数:

1
2
$ python func_cmdline.py --a 1 --d '2021-01-01 00:00:00' '2021-02-01 00:00:00' '2021-03-01 00:00:00'
{'a': 1, 'b': False, 'd': [datetime.datetime(2021, 1, 1, 0, 0), datetime.datetime(2021, 2, 1, 0, 0), datetime.datetime(2021, 3, 1, 0, 0)], 'e': 1j, 'c': ()}

对各个参数随便传入一些可用的值:

1
2
$ python func_cmdline.py --a 1 -b -c 2 3 4 --d '2021-01-01 00:00:00' '2021-02-01 00:00:00' '2021-03-01 00:00:00' -e '5+6j'
{'a': 1, 'b': True, 'd': [datetime.datetime(2021, 1, 1, 0, 0), datetime.datetime(2021, 2, 1, 0, 0), datetime.datetime(2021, 3, 1, 0, 0)], 'e': (5+6j), 'c': (2.0, 3.0, 4.0)}
分享

码中人
作者
Web Developer