#| default_exp py2pyi

#| export
import ast, sys, inspect, re, os, importlib.util, importlib.machinery

from ast import parse, unparse
from inspect import signature, getsource
from fastcore.utils import *
from fastcore.meta import delegates

#| export
def imp_mod(module_path, package=None):
    "Import dynamically the module referenced in `fn`"
    module_path = str(module_path)
    module_name = os.path.splitext(os.path.basename(module_path))[0]
    spec = importlib.machinery.ModuleSpec(module_name, None, origin=module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader = importlib.machinery.SourceFileLoader(module_name, module_path)
    if package is not None: module.__package__ = package
    module.__file__ = os.path.abspath(module_path)
    spec.loader.exec_module(module)
    return module

fn = Path('test_py2pyi.py')

mod = imp_mod(fn)
a = mod.A()
a.h()

#| export
def _get_tree(mod):
    return parse(getsource(mod))

tree = _get_tree(mod)

#| export
@patch
def __repr__(self:ast.AST):
    return unparse(self)

@patch
def _repr_markdown_(self:ast.AST):
    return f"""```python
{self!r}
```"""

# for o in enumerate(tree.body): print(o)

node = tree.body[4]
node

#| export
functypes = (ast.FunctionDef,ast.AsyncFunctionDef)

isinstance(node, functypes)

#| export
def _deco_id(d:Union[ast.Name,ast.Attribute])->bool:
    "Get the id for AST node `d`"
    return d.id if isinstance(d, ast.Name) else d.func.id

def has_deco(node:Union[ast.FunctionDef,ast.AsyncFunctionDef], name:str)->bool:
    "Check if a function node `node` has a decorator named `name`"
    return any(_deco_id(d)==name for d in getattr(node, 'decorator_list', []))

nm = 'delegates'
has_deco(node, nm)

node = tree.body[5]
node

has_deco(node, nm)

def _proc_body   (node, mod): print('_proc_body', type(node))
def _proc_func   (node, mod): print('_proc_func', type(node))
def _proc_class  (node, mod): print('_proc_class', type(node))
def _proc_patched(node, mod): print('_proc_patched', type(node))

#| export
def _get_proc(node):
    if isinstance(node, ast.ClassDef): return _proc_class
    if not isinstance(node, functypes): return None
    if not has_deco(node, 'delegates'): return _proc_body
    if has_deco(node, 'patch'): return _proc_patched
    return _proc_func

#| export
def _proc_tree(tree, mod):
    for node in tree.body:
        proc = _get_proc(node)
        if proc: proc(node, mod)

#| export
def _proc_mod(mod):
    tree = _get_tree(mod)
    _proc_tree(tree, mod)
    return tree

_proc_mod(mod);

node.name

sym = getattr(mod, node.name)
sym

sig = signature(sym)
print(sig)

#| export
def sig2str(sig):
    s = str(sig)
    s = re.sub(r"<class '(.*?)'>", r'\1', s)
    s = re.sub(r"dynamic_module\.", "", s)
    return s

#| export
def ast_args(func):
    sig = signature(func)
    return ast.parse(f"def _{sig2str(sig)}: ...").body[0].args

newargs = ast_args(sym)
newargs

node.args

node.args = newargs
node

#| export
def _body_ellip(n: ast.AST):
    stidx = 1 if isinstance(n.body[0], ast.Expr) and isinstance(n.body[0].value, ast.Str) else 0
    n.body[stidx:] = [ast.Expr(ast.Constant(...))]

_body_ellip(node)
node

#| export
def _update_func(node, sym):
    """Replace the parameter list of the source code of a function `f` with a different signature.
    Replace the body of the function with just `pass`, and remove any decorators named 'delegates'"""
    node.args = ast_args(sym)
    _body_ellip(node)
    node.decorator_list = [d for d in node.decorator_list if _deco_id(d) != 'delegates']

tree = _get_tree(mod)
node = tree.body[5]
node

_update_func(node, sym)
node

#| export
def _proc_body(node, mod): _body_ellip(node)

#| export
def _proc_func(node, mod):
    sym = getattr(mod, node.name)
    _update_func(node, sym)

tree = _proc_mod(mod)
tree.body[5]

node = tree.body[9]
node

ann = node.args.args[0].annotation

if hasattr(ann, 'elts'): ann = ann.elts[0]

nm = ann.id
nm

cls = getattr(mod, nm)
sym = getattr(cls, node.name)

sig2str(signature(sym))

_update_func(node, sym)

node

#| export
def _proc_patched(node, mod):
    ann = node.args.args[0].annotation
    if hasattr(ann, 'elts'): ann = ann.elts[0]
    cls = getattr(mod, ann.id)
    sym = getattr(cls, node.name)
    _update_func(node, sym)

tree = _proc_mod(mod)
tree.body[9]

tree = _get_tree(mod)
node = tree.body[7]
node

node.body

#| export
def _proc_class(node, mod):
    cls = getattr(mod, node.name)
    _proc_tree(node, cls)

tree = _proc_mod(mod)
tree.body[7]

#| export
def create_pyi(fn, package=None):
    "Convert `fname.py` to `fname.pyi` by removing function bodies and expanding `delegates` kwargs"
    fn = Path(fn)
    mod = imp_mod(fn, package=package)
    tree = _proc_mod(mod)
    res = unparse(tree)
    fn.with_suffix('.pyi').write_text(res)

create_pyi(fn)

# fn = Path('/Users/jhoward/git/fastcore/fastcore/docments.py')
# create_pyi(fn, 'fastcore')

#| export
from fastcore.script import call_parse

#| export
@call_parse
def py2pyi(fname:str,  # The file name to convert
           package:str=None  # The parent package
          ):
    "Convert `fname.py` to `fname.pyi` by removing function bodies and expanding `delegates` kwargs"
    create_pyi(fname, package)

#| hide
import nbdev; nbdev.nbdev_export()