generated from fastai/nbdev_template
-
Notifications
You must be signed in to change notification settings - Fork 262
/
script.py
119 lines (103 loc) · 4.47 KB
/
script.py
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
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/08_script.ipynb (unless otherwise specified).
__all__ = ['store_true', 'store_false', 'bool_arg', 'clean_type_str', 'Param', 'anno_parser', 'args_from_prog',
'SCRIPT_INFO', 'call_parse']
# Cell
import inspect,functools,argparse,shutil
from .imports import *
from .utils import *
from .docments import docments
# Cell
def store_true():
"Placeholder to pass to `Param` for `store_true` action"
pass
# Cell
def store_false():
"Placeholder to pass to `Param` for `store_false` action"
pass
# Cell
def bool_arg(v):
"Use as `type` for `Param` to get `bool` behavior"
return str2bool(v)
# Cell
def clean_type_str(x:str):
x = str(x)
x = re.sub("(enum |class|function|__main__\.|\ at.*)", '', x)
x = re.sub("(<|>|'|\ )", '', x) # spl characters
return x
# Cell
class Param:
"A parameter in a function used in `anno_parser` or `call_parse`"
def __init__(self, help=None, type=None, opt=True, action=None, nargs=None, const=None,
choices=None, required=None, default=None):
if type in (store_true,bool): type,action,default=None,'store_true' ,False
if type==store_false: type,action,default=None,'store_false',True
if type and isinstance(type,typing.Type) and issubclass(type,enum.Enum) and not choices: choices=list(type)
store_attr()
def set_default(self, d):
if self.default is None:
if d==inspect.Parameter.empty: self.opt = False
else: self.default = d
if self.default is not None: self.help += f" (default: {self.default})"
@property
def pre(self): return '--' if self.opt else ''
@property
def kwargs(self): return {k:v for k,v in self.__dict__.items()
if v is not None and k!='opt' and k[0]!='_'}
def __repr__(self):
if not self.help and self.type is None: return ""
if not self.help and self.type is not None: return f"{clean_type_str(self.type)}"
if self.help and self.type is None: return f"<{self.help}>"
if self.help and self.type is not None: return f"{clean_type_str(self.type)} <{self.help}>"
# Cell
class _HelpFormatter(argparse.HelpFormatter):
def __init__(self, prog, indent_increment=2):
cols = shutil.get_terminal_size((120,30))[0]
super().__init__(prog, max_help_position=cols//2, width=cols, indent_increment=indent_increment)
def _expand_help(self, action): return self._get_help_string(action)
# Cell
def anno_parser(func, prog=None, from_name=False):
"Look at params (annotated with `Param`) in func and return an `ArgumentParser`"
p = argparse.ArgumentParser(description=func.__doc__, prog=prog, formatter_class=_HelpFormatter)
for k,v in docments(func, full=True, returns=False, eval_str=True).items():
param = v.anno
if not isinstance(param,Param): param = Param(v.docment, v.anno)
param.set_default(v.default)
p.add_argument(f"{param.pre}{k}", **param.kwargs)
p.add_argument(f"--pdb", help=argparse.SUPPRESS, action='store_true')
p.add_argument(f"--xtra", help=argparse.SUPPRESS, type=str)
return p
# Cell
def args_from_prog(func, prog):
"Extract args from `prog`"
if prog is None or '#' not in prog: return {}
if '##' in prog: _,prog = prog.split('##', 1)
progsp = prog.split("#")
args = {progsp[i]:progsp[i+1] for i in range(0, len(progsp), 2)}
annos = type_hints(func)
for k,v in args.items():
t = annos.get(k, Param()).type
if t: args[k] = t(v)
return args
# Cell
SCRIPT_INFO = SimpleNamespace(func=None)
# Cell
def call_parse(func):
"Decorator to create a simple CLI from `func` using `anno_parser`"
mod = inspect.getmodule(inspect.currentframe().f_back)
if not mod: return func
@functools.wraps(func)
def _f(*args, **kwargs):
mod = inspect.getmodule(inspect.currentframe().f_back)
if not mod: return func(*args, **kwargs)
if not SCRIPT_INFO.func and mod.__name__=="__main__": SCRIPT_INFO.func = func.__name__
if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)
p = anno_parser(func)
args = p.parse_args().__dict__
xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)
tfunc = trace(func) if args.pop('pdb', False) else func
tfunc(**merge(args, args_from_prog(func, xtra)))
if mod.__name__=="__main__":
setattr(mod, func.__name__, _f)
SCRIPT_INFO.func = func.__name__
return _f()
else: return _f