Python AST Transformer

P2M.py

import ast
import inspect
import types


class Apply(ast.NodeTransformer):

    def __init__(self):
        super().__init__()
        self._never_function = True

    def visit_FunctionDef(self, node):
        if self._never_function:
            self._never_function = False
            decorator_list = []
            meet_me = False
            for decorator in node.decorator_list:
                if meet_me:
                    decorator_list.append(decorator)
                if isinstance(decorator, ast.Attribute):
                    if decorator.value.id == type(self).__name__:
                        meet_me = True
            result = ast.FunctionDef(
                name=node.name,
                args=node.args,
                body=node.body,
                returns=node.returns,
                type_comment=node.type_comment,
                decorator_list=decorator_list,
            )
        else:
            return result
        return self.generic_visit(result)

    @classmethod
    def apply(cls, func):
        source = inspect.getsource(func)
        tree = ast.parse(source, mode="exec")
        new_tree = cls().visit(tree)
        new_source = ast.unparse(ast.fix_missing_locations(new_tree))
        new_code = compile(
            ast.fix_missing_locations(new_tree),
            inspect.getfile(func),
            "exec",
        )
        new_func = types.FunctionType(new_code.co_consts[0], func.__globals__)
        return new_func


class P2M(Apply):

    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            result = ast.BinOp(left=node.left, op=ast.Mult(), right=node.right)
        elif isinstance(node.op, ast.Mult):
            result = ast.BinOp(left=node.left, op=ast.Add(), right=node.right)
        else:
            result = node
        return self.generic_visit(result)


# Demo


@P2M.apply
def add(a, b):
    return a + b


print(add(2, 3))


@P2M.apply
def func(a, b, c, d):
    return (a + 2) * b - (c + d)


print(func(2, 3, 4, 5))