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
|
__author__ = "ChenyangGao <https://chenyanggao.github.io/>" __version__ = (0, 1, 1) __all__ = ["func_cmdline"]
import inspect
from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter from functools import partial from textwrap import dedent, indent from typing import 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, ) params = inspect.signature(fn).parameters if parser is None: parser = ArgumentParser( formatter_class=RawDescriptionHelpFormatter) 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: return parser.parse_args(argv) def run(argv: Optional[list[str]] = None): "解析命令行参数,默认使用" 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__": @func_cmdline(helpinfo_map=dict( a="参数a", b="参数b", c="参数c", d="参数d", e="参数e", f="参数f", )) def foo(a: int, /, b: bool, c: bool = True, *d: int, e: list[int], f: int = 4): "Example for function as command line." print(a, b, c, d, e, f)
foo.run()
|