How to pass a user defined argument to setuptools in order to set a flag that changes compilation macro

I have some large setup.py file that compiles several CUDA files, something like (VERY INCOMPETE, I can provide more info if its relevant):

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=[("MACRO_I_WANT", None)],

    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/"],
)

# etc

setup(
    name="-",
    version="-",
    author="-",
    packages=find_packages(),
    include_package_data=True,
    data_files=[("data", ["../data/somefile.file"])],
    ext_modules=[foo1, foo2, gpuUtils_ext, foo3],   # I have many
    py_modules=["foo.py"],
    cmdclass={"build_ext": BuildExtension},
    install_requires=["Cython", "matplotlib", "numpy", "scipy", "tqdm"],
    license_files=("LICENSE",),
    license="BSD 3-Clause",
    # since the package has c code, the egg cannot be zipped
    zip_safe=False,
)

The file gpuUtils.cu, that is being compiled in this setup.py has a macro MACRO_I_WANT that its defined here and the inside the file there is a #ifdef to disable a piece of code.

I would like to change setuptools such that the user can provide a flag for this macro, e.g. python setup.py install would not define the macro, but python setup.py intall -define-macro would define it.

As far as I can see/test, the general option of distutils: How to pass a user defined parameter to setup.py? does not work, because by the time InstallCommand is called, my Extensions have already been defined and passed to setup.

Is this doable? How can I do it? Is this the right way of approaching it?

Solution

Ended up grabbing the system argument and poping it out:

import sys
# Macros for compilation
define_macros=[("SOME_MACRO", None)]
my_flag=False
if len(sys.argv)>1:
    if "-define-macro " in sys.argv[1:] :
        my_flag=True
        sys.argv.pop(sys.argv.index("-define-macro"))
    else:
        raise ValueError("flag not understood, only -define-macro accepted")
 
if my_flag:
    define_macros.append(("MACRO_I_WANT",None))

Then, in the extensions:

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=define_macros, # changed this

    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/"],
)