Function to transform each node of `tree` matching types in `names` by passing them to `func`

tfm_ast.py

import ast
from fastcore.utils import *
from fastcore.utils import _get_class

def tfm_ast(tree:ast.AST, func:callable, *names, **kwargs):
    "Transform each node of `tree` by passing it to `func`"
    def f(self, node):
        res = func(node, **kwargs)
        return res if res else node

    d = {'visit_'+nm:f for nm in names}
    res = _get_class('TmpVisitor', sup=ast.NodeTransformer, **d)().visit(tree)
    ast.fix_missing_locations(res)
    return res