#| 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()
1
#| 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
def f(a: int, b: str='a') -> str:
"""I am f"""
return 1
#| export
functypes = (ast.FunctionDef,ast.AsyncFunctionDef)
isinstance(node, functypes)
True
#| 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)
False
node = tree.body[5]
node
@delegates(f)
def g(c, d: X, **kwargs) -> str:
"""I am g"""
return 2
has_deco(node, nm)
True
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);
_proc_class <class 'ast.ClassDef'> _proc_body <class 'ast.FunctionDef'> _proc_func <class 'ast.FunctionDef'> _proc_body <class 'ast.FunctionDef'> _proc_class <class 'ast.ClassDef'> _proc_class <class 'ast.ClassDef'> _proc_patched <class 'ast.FunctionDef'> _proc_patched <class 'ast.FunctionDef'> _proc_body <class 'ast.FunctionDef'>
node.name
'g'
sym = getattr(mod, node.name)
sym
<function test_py2pyi.g(c, d: test_py2pyi.X, *, b: str = 'a') -> str>
sig = signature(sym)
print(sig)
(c, d: test_py2pyi.X, *, b: str = 'a') -> str
#| 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
c, d: test_py2pyi.X, *, b: str='a'
node.args
c, d: X, **kwargs
node.args = newargs
node
@delegates(f)
def g(c, d: test_py2pyi.X, *, b: str='a') -> str:
"""I am g"""
return 2
#| 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
@delegates(f)
def g(c, d: test_py2pyi.X, *, b: str='a') -> str:
"""I am g"""
...
#| 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
@delegates(f)
def g(c, d: X, **kwargs) -> str:
"""I am g"""
return 2
_update_func(node, sym)
node
def g(c, d: test_py2pyi.X, *, b: str='a') -> str:
"""I am g"""
...
#| 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]
_proc_class <class 'ast.ClassDef'> _proc_class <class 'ast.ClassDef'> _proc_class <class 'ast.ClassDef'> _proc_patched <class 'ast.FunctionDef'> _proc_patched <class 'ast.FunctionDef'>
def g(c, d: test_py2pyi.X, *, b: str='a') -> str:
"""I am g"""
...
node = tree.body[9]
node
@patch
@delegates(j)
def k(self: (A, B), b: bool=False, **kwargs):
return 1
ann = node.args.args[0].annotation
if hasattr(ann, 'elts'): ann = ann.elts[0]
nm = ann.id
nm
'A'
cls = getattr(mod, nm)
sym = getattr(cls, node.name)
sig2str(signature(sym))
"(self: (test_py2pyi.A, test_py2pyi.B), b: bool = False, *, d: str = 'a')"
_update_func(node, sym)
node
@patch
def k(self: (test_py2pyi.A, test_py2pyi.B), b: bool=False, *, d: str='a'):
...
#| 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]
_proc_class <class 'ast.ClassDef'> _proc_class <class 'ast.ClassDef'> _proc_class <class 'ast.ClassDef'>
@patch
def k(self: (test_py2pyi.A, test_py2pyi.B), b: bool=False, *, d: str='a'):
...
tree = _get_tree(mod)
node = tree.body[7]
node
class A:
@delegates(j)
def h(self, b: bool=False, **kwargs):
a = 1
return a
node.body
[@delegates(j) def h(self, b: bool=False, **kwargs): a = 1 return a]
#| export
def _proc_class(node, mod):
cls = getattr(mod, node.name)
_proc_tree(node, cls)
tree = _proc_mod(mod)
tree.body[7]
class A:
def h(self, b: bool=False, *, d: str='a'):
...
#| 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()