script.py

"""A fast way to turn your python function into a script.

Docs: https://fastcore.fast.ai/script.html.md"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_script.ipynb.

# %% auto #0
__all__ = ['store_true', 'store_false', 'bool_arg', 'anno_parser', 'args_from_prog', 'set_ctx', 'call_parse', 'is_cli']

# %% ../nbs/06_script.ipynb #8a36db98
import inspect,argparse,shutil,types,asyncio

from functools import wraps,partial
from .imports import *
from .utils import *
from .docments import docments, ann_parts
from typing import get_origin, get_args, Union

# %% ../nbs/06_script.ipynb #5bf7ac6c
def store_true():
    "Placeholder annotation type for a `store_true` argparse action"
    pass

# %% ../nbs/06_script.ipynb #e2798370
def store_false():
    "Placeholder annotation type for a `store_false` argparse action"
    pass

# %% ../nbs/06_script.ipynb #69d525bc
def bool_arg(v):
    "Annotation type giving `bool` behavior for CLI args"
    return str2bool(v)

# %% ../nbs/06_script.ipynb #6785a783
def _arg_kw(k, anno, doc, default, extra):
    "CLI arg name and `add_argument` kwargs for param `k`"
    extra = dict(extra)
    action,d = extra.get('action'),extra.pop('default', default)
    opt,negated = extra.pop('opt', None),False
    if   anno==store_true:  action,d,anno = 'store_true',False,None
    elif anno==store_false: action,d,anno = 'store_false',True,None
    elif anno==bool and action is None:
        anno = None
        if d is True: action,negated = 'store_false',True
        else: action,d = 'store_true',False
    if action=='version':
        if 'version' not in extra and d is not inspect.Parameter.empty: extra['version'] = d
        return f'--{k}', {'help':doc or '', **extra}
    kw = {}
    if action and 'action' not in extra: kw['action'] = action
    if anno is not None:
        kw['type'] = anno
        if isinstance(anno,type) and issubclass(anno,enum.Enum): kw['choices'] = list(anno)
    if d is not inspect.Parameter.empty and d is not None: kw['default'] = d
    if opt is None: opt = d is not inspect.Parameter.empty
    kw['help'] = (doc or '') + (f" (default: {kw['default']})" if 'default' in kw else '')
    if negated: kw['dest'] = k
    return f"{'--' if opt else ''}{f'no-{k}' if negated else k}", {**kw, **extra}

# %% ../nbs/06_script.ipynb #cceb4486
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)

# %% ../nbs/06_script.ipynb #2259c093
def _is_union(t): return get_origin(t) in (Union, types.UnionType) if hasattr(types, 'UnionType') else get_origin(t) is Union

def _union_parser(types):
    "Return a parser that tries each type in sequence"
    def _parse(v):
        for t in types:
            if t is type(None): continue
            try: return t(v)
            except: pass
        raise ValueError(f"Could not parse {v!r} as any of {types}")
    return _parse

def _union_type(t):
    "Get parser for Union types, or None if not a Union"
    if not _is_union(t): return None
    return _union_parser(get_args(t))

# %% ../nbs/06_script.ipynb #5e5bea67
def anno_parser(func, prog:str=None):
    "Look at params (with type/docments/`Annotated` annotations) 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():
        anno,meta = ann_parts(v.anno)
        extra = next((o for o in meta if isinstance(o,dict)), {})
        anno = _union_type(anno) or anno
        name,kw = _arg_kw(k, anno, v.docment, v.default, extra)
        p.add_argument(name, **kw)
    p.add_argument(f"--pdb", help=argparse.SUPPRESS, action='store_true')
    p.add_argument(f"--xtra", help=argparse.SUPPRESS, type=str)
    return p

# %% ../nbs/06_script.ipynb #75e8a419
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,meta = ann_parts(annos.get(k))
        extra = next((o for o in meta if isinstance(o,dict)), {})
        if t in (bool, bool_arg, store_true, store_false) or extra.get('action') in ('store_true','store_false'): t = str2bool
        if t: args[k] = t(v)
    return args

# %% ../nbs/06_script.ipynb #42c8e85f
from contextvars import ContextVar
from contextlib import contextmanager

# %% ../nbs/06_script.ipynb #e4537112
@contextmanager
def set_ctx(cv, val=True):
    token = cv.set(val)
    try: yield
    finally: cv.reset(token)

# %% ../nbs/06_script.ipynb #fc816498
def _is_script_run(frame):
    "True if `frame` is the top-level body of a file being run directly (`python foo.py`, `python -m foo`, or `%run foo.py`)"
    if not frame or frame.f_locals is not frame.f_globals: return False
    g = frame.f_globals
    if g.get('__name__')!='__main__' or not g.get('__file__'): return False
    return Path(g['__file__']).resolve()==Path(frame.f_code.co_filename).resolve()

# %% ../nbs/06_script.ipynb #dee5e259
_cli_func = ContextVar('_cli_func', default=None)

def _run_cli(func, nested):
    "Parse `sys.argv` with `anno_parser` and call `func` with the result"
    with set_ctx(_cli_func, func):
        if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)
        p = anno_parser(func)
        if nested: args, sys.argv[1:] = p.parse_known_args()
        else: args = p.parse_args()
        args = args.__dict__
        xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)
        tfunc = trace(func) if args.pop('pdb', False) else func
        res = tfunc(**merge(args, args_from_prog(func, xtra)))
        return asyncio.run(res) if inspect.isawaitable(res) else res

def call_parse(func=None, nested=False):
    "Decorator to create a simple CLI from `func` using `anno_parser`"
    if func is None: return partial(call_parse, nested=nested)
    @wraps(func)
    def _f(*args, **kwargs):
        if args or kwargs or _cli_func.get() is not None: return func(*args, **kwargs)
        if _is_script_run(inspect.currentframe().f_back): return _run_cli(func, nested)
        with set_ctx(_cli_func, False): return func(*args, **kwargs)

    frame = inspect.currentframe().f_back
    if _is_script_run(frame):
        frame.f_globals[func.__name__] = _f
        return _run_cli(func, nested)
    else: return _f

# %% ../nbs/06_script.ipynb #f4961093
def is_cli(func=None):
    "True if a `call_parse` CLI run is in progress, optionally checking that `func` is the function being run"
    f = _cli_func.get()
    return bool(f) if func is None else f is getattr(func, '__wrapped__', func)