如何在导入Python模块时传入参数，实现动态选择导入numpy或cupy包？

Python的import语句本身不支持直接传入参数来控制模块的导入行为，但我们可以通过封装逻辑的方式实现你想要的效果。下面是几种可行的方案，你可以根据需求选择：

方案一：使用全局初始化函数（简单直接） #

这种方式通过在模块中定义一个初始化函数，导入模块后调用该函数并传入CUDA参数，完成动态导入和配置。

my_module.py #

# 定义全局变量存储选中的包
pkg = None

def setup(cuda: bool):
    global pkg
    # 根据参数选择导入的包
    if cuda:
        import cupy
        pkg = cupy
    else:
        import numpy
        pkg = numpy
    # 打印包名称
    print(pkg.__name__)

main.py #

import my_module

# 传入CUDA参数配置模块
my_module.setup(CUDA=True)  # 控制台输出 "cupy"
# 或者
# my_module.setup(CUDA=False) # 控制台输出 "numpy"

# 后续可以通过my_module.pkg调用对应的库功能
# 示例：创建数组
# arr = my_module.pkg.array([1,2,3])

方案二：使用类封装（更优雅，避免全局变量） #

如果想避免全局变量的使用，可以用类来封装后端逻辑，通过实例化类时传入参数来控制导入。

my_module.py #

class Backend:
    def __init__(self, cuda: bool):
        # 初始化时根据参数导入对应包
        if cuda:
            import cupy
            self.pkg = cupy
        else:
            import numpy
            self.pkg = numpy
        # 打印包名称
        print(self.pkg.__name__)

main.py #

from my_module import Backend

# 实例化时传入CUDA参数
backend = Backend(CUDA=True)  # 输出 "cupy"
# 或者
# backend = Backend(CUDA=False) # 输出 "numpy"

# 使用backend.pkg调用库功能
# arr = backend.pkg.array([1,2,3])

方案三：通过环境变量控制（适合导入阶段自动配置） #

如果希望在import模块的同时完成配置，可以借助环境变量，在导入前设置环境变量，模块读取该变量来决定导入哪个包。

main.py #

import os

# 设置环境变量控制CUDA开关
os.environ['MY_MODULE_CUDA'] = 'True'  # 设为'False'则导入numpy
# os.environ['MY_MODULE_CUDA'] = 'False'

import my_module

my_module.py #

import os

# 读取环境变量，默认关闭CUDA
cuda_enabled = os.getenv('MY_MODULE_CUDA', 'False').lower() == 'true'

if cuda_enabled:
    import cupy as pkg
else:
    import numpy as pkg

print(pkg.__name__)

为什么原思路不可行？ #

你最初的想法是在模块顶层直接根据参数导入包，但模块的顶层代码会在import时立即执行，而import语句本身不支持传入参数，所以无法在导入时动态传递CUDA参数。必须将导入逻辑封装到函数、类等可调用结构中，才能在导入后传入参数控制行为。

内容的提问来源于stack exchange，提问作者Sam-gege