如何在Python包安装时一次性计算并保存耗时常量？

问题 #

使用setup.py安装Python模块时，子模块module_name.submodule中存在一项耗时极长的一次性计算，生成的CONST_VARS常量不会变更，但每次加载模块时都会重复执行该计算，导致启动耗时过长。需求是仅在包安装阶段计算并保存这些常量，且不希望手动将常量写入脚本或文件。

给出的代码如下：

setup.py #

from setuptools import setup, find_packages
import module_name

setup(
      name = 'module_name',
      ...
      packages = find_packages()
     )

module_name/submodule.py #

def func1():
    ...

CONST_VARS = []
... calculate here (takes a long time)...

解决方案 #

完全可以通过在setup.py执行阶段生成常量文件的方式实现需求，具体步骤如下：

1. 抽离计算逻辑为独立函数 #

将submodule.py中的耗时计算逻辑抽成可调用的函数，建议放到module_name/utils.py中（也可保留在submodule.py内，抽离后结构更清晰）：

# module_name/utils.py
def get_const_vars():
    # 原耗时计算逻辑
    CONST_VARS = []
    # ... 这里是你的长耗时计算代码 ...
    return CONST_VARS

2. 修改setup.py，在安装阶段生成常量文件 #

修改setup.py，在执行setup()前调用计算函数，将结果写入子模块目录下的专用常量文件（比如_constants.py）：

from setuptools import setup, find_packages
import os
import sys

# 临时添加当前目录到Python路径，确保能导入模块内的函数
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from module_name.utils import get_const_vars

# 生成常量文件路径
constants_file = os.path.join('module_name', '_constants.py')
# 执行计算并写入文件
const_values = get_const_vars()
with open(constants_file, 'w', encoding='utf-8') as f:
    f.write(f'CONST_VARS = {repr(const_values)}\n')

setup(
    name='module_name',
    # ... 其他原有参数 ...
    packages=find_packages(),
    # 确保生成的常量文件被打包进安装包
    package_data={'module_name': ['_constants.py']},
)

注意：如果计算依赖第三方包，需确保setup.py的install_requires已包含这些依赖，否则安装阶段执行计算会报错。

3. 修改submodule.py，导入预生成的常量 #

删掉submodule.py中原有的耗时计算代码，改为从生成的_constants.py导入常量：

# module_name/submodule.py
def func1():
    ...

# 替换原计算代码，直接导入预生成的常量
from ._constants import CONST_VARS

4. 可选：避免开发模式重复生成（按需） #

如果使用python setup.py develop或pip install -e .进行开发，可在setup.py中添加判断，仅当_constants.py不存在时才生成，避免每次开发安装都重复计算：

# 在生成常量文件前添加判断
if not os.path.exists(constants_file):
    const_values = get_const_vars()
    with open(constants_file, 'w', encoding='utf-8') as f:
        f.write(f'CONST_VARS = {repr(const_values)}\n')

验证 #

执行python setup.py install或pip install .安装包后，查看site-packages/module_name/_constants.py会发现已写入计算好的CONST_VARS值。后续加载module_name.submodule时，会直接导入该常量，不再执行耗时计算。

内容的提问来源于stack exchange，提问作者sodiumnitrate