如何在 Python 中做到类似 #ifdef DEBUG

类似

#ifndef NDEBUG
do_something(...)
#else
do_otherthing(...)
#endif

logging 模块

如果有无 NDEBUG，只在于是否输出一些信息，那么可以使用 logging 模块，这是一个用于记录和管理日志信息的标准库，通过配置日志级别，可以控制不同等级的日志的输出。

import logging
logging.basicConfig(level=logging.INFO, filename='log', filemode='w',
                    format='%(asctime)s - %(levelname)s - %(message)s')
logging.debug('这是一个 debug 级别的日志消息')
logging.info('这是一个 info 级别的日志消息')
logging.warning('这是一个 warning 级别的日志消息')
logging.error('这是一个 error 级别的日志消息')

以上代码会将 INFO 等级及以上的消息输出到 log 文件中，即 logging.debug 这句不会输出。

以上是基本使用方法，还可以根据实际需求进行更高级的配置，例如添加日志处理器、设置日志过滤器等。更详细的文档可以参考 Python 官方文档或其他教程。

全局变量

DEBUG = True

if DEBUG:
    print('[DEBUG]')

简单易懂，但因为 Python 没有真正的常变量，实际上和 #ifndef NDEBUG 还是有一点差距的。可以使用 __debug__ 取代这一全局变量。

Python 会对 __debug__ 进行优化，if __debug__ 等同于 if False 或 if True，取决于运行脚本时是否有加上 -O 参数。不同于使用全局变量 DEBUG，if False if True 会被优化，可以使用 dis 模块反汇编观察 Python 的字节码。

文章后面 DEBUG __debug__ 都是可以互换的，区别只是一个是全局变量（可以放在一个包中引用），一个是 Python 提供的常量；并且后续提供的几种方法有些是可以结合的。

装饰器

假设所有在条件执行的语句都可以被封装在函数中，那么可以写一个装饰器来取代全局变量 DEBUG 令语言更加整洁。

写一个判断数字是否是素数的函数 is_prime(n)。

from typing import Generator
from dis import dis

def debug_decorator(func):
    def wrapper(*args, **kwargs):
        if __debug__:
            print(f"[DEBUG {func.__name__}]")
            return func(*args, **kwargs)
        return None
    return wrapper


def is_prime(n: int) -> bool:
    Gen = Generator[int, None, None]

    def prime_filter(g: Gen) -> Gen:
        x = next(g)
        yield x
        yield from prime_filter(filter(lambda y: y % x != 0, g))

    return n in prime_filter(iter(range(2, n + 1)))


@debug_decorator
def f():
    assert is_prime(47)
    print(47)

print(dis(f))

使用 python -u main.py 运行

  9           0 LOAD_GLOBAL              0 (print)
              2 LOAD_CONST               1 ('[DEBUG ')
              4 LOAD_DEREF               0 (func)
              6 LOAD_ATTR                1 (__name__)
              8 FORMAT_VALUE             0
             10 LOAD_CONST               2 (']')
             12 BUILD_STRING             3
             14 CALL_FUNCTION            1
             16 POP_TOP

 10          18 LOAD_DEREF               0 (func)
             20 LOAD_FAST                0 (args)
             22 LOAD_FAST                1 (kwargs)
             24 CALL_FUNCTION_EX         1
             26 RETURN_VALUE
None

使用 python -u -O main.py 运行（加上 -O），函数变为单纯返回 None 的函数。

 11           0 LOAD_CONST               0 (None)
              2 RETURN_VALUE
None

ast 模块

注意 Python 版本。

还可以像真正的预编译器一样尝试对源文件进行修改。ast 可以将 Python 源文件转化为抽象语法树（AST），我们不直接对源文件进行修改，而是在抽象语法树上进行修改，再转化为源文件。

假设现在有源文件如下，要求删除 with debug_block 语句块。

class DebugBlock:
    def __enter__(self): pass
    def __exit__(self, exc_type, exc_value, exc_traceback): pass

debug_block = DebugBlock()

with debug_block:
    print("debug")

print("info")

用到 ast.NodeTransformer就可以在 AST 上删除 with debug_block 所代表的的语句块。需要实现 visit_With，传入的参数 node 就是一个 with 语句块，在 visit_With 中可以修改 with 语句块，选择 id 为 debug_block 的 with 语句块并赋值空语句，就将带有 debug_block 的语句块删除了。

import ast

class DebugTransformer(ast.NodeTransformer):
    def visit_With(self, node: ast.FunctionDef):
        # self.generic_visit(node)
        for item in node.items:
            try:
                id = item.context_expr.id
                if id == "debug_block":
                    node = ast.parse("")
            except AttributeError:
                continue
        return node

对于单一文件，可以直接使用 exec 运行，比如对于源文件 main.py 可以另造一个 run.py，进行读取 main.py 文件、转换、运行，三步。

# main.py
with debug_block:
    print("debug")
print("info")

# run.py
with open("main.py", "r") as source:
    # source code to AST
    tree = ast.parse(source.read())
# transform: remove with debug_block: ...
node = DebugTransformer().visit(tree)
# AST to source code and exec
exec(ast.unparse(node))

这样运行 run.py 则只会输出 info。

对于处于一个文件夹下的多个源文件则遍历转换即可。

def transform(src_dir=os.getcwd()):
    dst_dir = f'{src_dir}.pre'
    shutil.copytree(src_dir, dst_dir)

    for root, _, files in os.walk(dst_dir):
        for f in files:
            path = os.path.join(root, f)
            with open(path, 'r') as source:
                code = source.read()
            tree = ast.parse(code)
            tree = DebugTransformer().visit(tree)
            code = ast.unparse(node)
            with open(path, 'w') as source:
                source.write(code)

看上去这种方法是最复杂的，但这正是因为 ast 模块比较底层，深入到了语法层面。ast 生成的抽象语法树只关系语法层面，不考虑语义（因此前面的 debug_block 只是一个符号 Python 语法的标识符，不需要真的是一个有效值）；所以使用 ast.NodeTransformer 可以转换语法，自定义很多操作，甚至可以几乎变成另一种语言。

不过 ast 库本身并不是很完善的一个库，更多地是用在语法研究中，这里也只是简单介绍一下。

第三方库 pypreprocessor

让你真的使用 C 语言的宏表达式。pypreprocessor 本身的使用也比较简单，毕竟说是库，但其实只有一个源文件。

pip3 -m pip install pypreprocessor

注意 # 和 ifdef else 等之间不能有空格，也就是 # if 不会被识别。不少自动格式化软件，比如 autopep8 会自动添加空格。

from pypreprocessor import pypreprocessor

pypreprocessor.parse()

#define DEBUG

#ifdef DEBUG
print("debug output")
#else
print("not debug output")
#endif

pypreprocessor 实际上是经过三步：预编译、运行预编译后的文件、删除文件，这和 ast 模块中所做的事是类似的。不过 pypreprocessor 是在源文件上作变动，而 ast 是在抽象语法树上做变动，也许可以将两者结合起来获得一个更好的预编译库。