如何向setuptools传递自定义参数以控制编译宏的定义

这个问题我之前也碰到过，确实因为Extension实例是在setup()函数调用前就创建好了，所以常规的distutils命令扩展方法会失效——等InstallCommand执行的时候，扩展模块的配置已经固定了。不过我们可以通过提前解析命令行参数的方式动态调整编译宏，完美解决你的需求，具体步骤如下：

1. 提前解析自定义命令行参数 #

在setup.py的最开头，先读取并处理用户传入的自定义标志，同时把这个标志从命令行参数列表中移除（避免setuptools因为识别不了这个参数而报错）：

import sys

# 检查用户是否传入了自定义宏开关参数
enable_macro_i_want = False
# 这里用--define-macro作为触发标志，你也可以改成更明确的名字比如--enable-macro-i-want
if '--define-macro' in sys.argv:
    enable_macro_i_want = True
    sys.argv.remove('--define-macro')  # 移除这个参数，防止setuptools报错

2. 动态构建Extension的define_macros参数 #

根据上面解析到的标志，决定是否添加MACRO_I_WANT到编译宏列表中：

# 先初始化编译宏列表
gpu_utils_defines = []
if enable_macro_i_want:
    gpu_utils_defines.append(("MACRO_I_WANT", None))

# 然后用这个动态生成的列表创建Extension
gpuUtils_ext = Extension(
    "_gpuUtils",
    sources=include_headers(
        [
            "gpuUtils.cu",
            "python/utilities/cuda_interface/_gpuUtils.pxd",
            "python/utilities/cuda_interface/_gpuUtils.pyx",
        ],
        sdist=sys.argv[1] == "sdist",
    ),
    define_macros=gpu_utils_defines,  # 这里替换成动态生成的列表
    library_dirs=[CUDA["lib64"]],
    libraries=["cudart"],
    language="c++",
    runtime_library_dirs=[CUDA["lib64"]] if not IS_WINDOWS else None,
    include_dirs=[NUMPY_INCLUDE, CUDA["include"], "./CUDA/"],
)

3. 剩下的setup()代码保持不变 #

你的setup()函数不需要做任何修改，直接使用上面创建的gpuUtils_ext即可。

使用方式 #

• 默认安装（不定义宏）：

  python setup.py install
  

• 开启宏定义安装：

  python setup.py install --define-macro
  

补充：用pip安装时的传递方式 #

如果用户是通过pip安装你的包，需要用--install-option来传递这个自定义参数：

pip install . --install-option="--define-macro"

为什么这个方法可行？ #

因为我们在setup()函数执行之前就完成了命令行参数的解析和Extension的配置，这样当setuptools开始处理安装流程时，Extension的编译宏已经根据用户的选择设置好了，完美避开了“Extension提前实例化导致无法修改”的问题。

如果之后需要控制更多编译宏，只需要扩展参数解析的逻辑即可，比如支持--define-macro=MACRO1 --define-macro=MACRO2=VALUE这种更灵活的形式，不过对于你当前的需求，上面的方案已经足够简单高效了。

备注：内容来源于stack exchange，提问作者Ander Biguri