Automatically decorating all top-level functions in Python using the ast module.

autodecorator.py

import ast
import sys

import astunparse


class Transformer(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        if node.col_offset == 0:
            decorator = ast.Name(id="log_call", ctx=ast.Load())
            return ast.FunctionDef(
                node.name,
                node.args,
                node.body,
                [decorator, *node.decorator_list],
                node.returns,
                node.type_comment,
            )
        return node


if __name__ == "__main__":
    with open(sys.argv[1], "r") as f:
        tree = ast.parse(f.read())
    transformed_tree = Transformer().visit(tree)
    transformed_source = astunparse.unparse(transformed_tree)
    sys.stdout.write(transformed_source)

input.py

import functools
import logging
from functools import lru_cache

log = logging.getLogger(__name__)

RANDOM_CONSTANT_2 = 1


@functools.lru_cache
def public_func_1(foo, bar):
    """hello there"""
    log.debug("test %d", 1)
    return 1 + 2 + 3


@lru_cache
def public_func_2(foo, bar):
    log.debug("test %d", 1)

    def inner_function(foo, bar):
        log.debug("test %d", 1)
        return 1 + 2 + 3

    return inner_function(foo, bar)


def public_func_3(foo, bar):
    """hej"""
    log.debug("test %d", 1)
    return 1 + 2 + 3


def _private_func(foo, bar):
    """hej"""
    log.debug("test %d", 1)
    return 1 + 2 + 3


RANDOM_CONSTANT_2 = 2

output.py

import functools
import logging
from functools import lru_cache

log = logging.getLogger(__name__)
RANDOM_CONSTANT_2 = 1


@log_call
@functools.lru_cache
def public_func_1(foo, bar):
    "hello there"
    log.debug("test %d", 1)
    return (1 + 2) + 3


@log_call
@lru_cache
def public_func_2(foo, bar):
    log.debug("test %d", 1)

    def inner_function(foo, bar):
        log.debug("test %d", 1)
        return (1 + 2) + 3

    return inner_function(foo, bar)


@log_call
def public_func_3(foo, bar):
    "hej"
    log.debug("test %d", 1)
    return (1 + 2) + 3


@log_call
def _private_func(foo, bar):
    "hej"
    log.debug("test %d", 1)
    return (1 + 2) + 3


RANDOM_CONSTANT_2 = 2