Python实现自动去除Debug代码的终极方案

# 误删
print = my_print
print("hello")   # 不该删

text = "print(x)"  # 字符串

import ast
import astor


DEBUG_FUNC_NAMES = {
    "print",
    "pprint",
    "debug",
}

LOGGING_METHODS = {
    "debug",
    "info",
}


class RemoveDebugTransformer(ast.NodeTransformer):
    def visit_Expr(self, node):
        """
        处理：
        - print(...)
        - logging.debug(...)
        - logger.debug(...)
        """
        call = node.value
        if not isinstance(call, ast.Call):
            return node

        func = call.func

        # print(...)
        if isinstance(func, ast.Name):
            if func.id in DEBUG_FUNC_NAMES:
                return None

        # logging.debug(...) / logger.debug(...)
        if isinstance(func, ast.Attribute):
            if func.attr in LOGGING_METHODS:
                return None

        return node

    def visit_If(self, node):
        """
        处理：
        if DEBUG:
            ...
        """
        # if DEBUG:
        if isinstance(node.test, ast.Name) and node.test.id == "DEBUG":
            return None

        return self.generic_visit(node)

def remove_debug_code(code: str) -> str:
    tree = ast.parse(code)

    transformer = RemoveDebugTransformer()
    tree = transformer.visit(tree)
    ast.fix_missing_locations(tree)

    return astor.to_source(tree)

import logging

DEBUG = True

print("hello")

logging.debug("debug log")
logging.info("info log")

logger.debug("logger debug")

x = 10

if DEBUG:
    print("only debug")

print("done")

code = """
import logging

DEBUG = True

def foo(x):
    print("foo x =", x)
    logging.debug("debug foo")
    logging.info("info foo")

    if DEBUG:
        print("only in debug")

    return x * 2


print("program start")
result = foo(10)
print("result =", result)
"""
new_code = remove_debug_code(code)
print(new_code)

import logging

x = 10

print("done")

if os.getenv("ENV") == "prod":
    code = remove_debug_code(code)

LOGGING_METHODS = {"debug", "info"}

def visit_Assert(self, node):
    return None

from pathlib import Path

for file in Path("src").rglob("*.py"):
    code = file.read_text(encoding="utf-8")
    new_code = remove_debug_code(code)
    file.write_text(new_code, encoding="utf-8")