如何向setuptools(setup.py)自动生成的启动脚本添加自定义导入语句

方案1：入口代码注入__main__命名空间（最推荐，零侵入兼容全版本） #

根本不需要修改setuptools的脚本生成逻辑，Pickle/PyTorch反序列化时找__main__下的属性，本质是查sys.modules['__main__']模块的全局命名空间，你只需要在入口函数执行的最开头，把需要的类手动注入进去就行，完全赶在模型加载前完成依赖注册。
在你的入口文件scripts/trans.py的main函数最顶部添加如下代码：

import sys
from casepunc.utils import WordTokenizer

# 把依赖类注入到__main__模块全局域
sys.modules["__main__"].WordTokenizer = WordTokenizer

这个方案不受setuptools版本、启动脚本模板变动的影响，跨平台兼容性和默认打包逻辑完全一致，不需要改任何打包配置。

方案2：重写Setuptools脚本生成逻辑（硬改生成的启动脚本） #

如果你一定要让导入语句直接出现在生成的启动脚本里，可以自定义脚本写入器，替换默认的生成模板。在setup.py中添加如下逻辑：

from setuptools import setup
from setuptools.command.easy_install import ScriptWriter

origin_get_args = ScriptWriter.get_args
def patched_script_writer(dist, exe_path, gui=False):
    for content, mode in origin_get_args(dist, exe_path, gui):
        # 在原有导入行后插入你需要的导入语句
        patched_content = content.replace(
            "from scripts.trans import main\n",
            "from scripts.trans import main\nfrom casepunc.utils import WordTokenizer\n"
        )
        yield patched_content, mode

ScriptWriter.get_args = staticmethod(patched_script_writer)

# 原有setup配置保持不变
setup(
    # 其余配置（name/version/依赖等）正常填写
    entry_points={
        "console_scripts": [
            # 你的命令配置和之前一致
            "your-cli-name = scripts.trans:main"
        ]
    }
)

注意：这个方案依赖setuptools内部的ScriptWriter实现和默认脚本模板，大版本升级后如果默认模板内容变了，字符串替换逻辑可能失效，需要同步调整。

方案3：使用固定自定义脚本替代自动生成 #

你可以完全放弃setuptools的自动脚本生成，自己写好固定内容的启动脚本，打包时直接指定脚本路径。

• 先在项目下新建启动脚本，比如路径为bin/your-cli-name，内容就是你预期的最终脚本内容：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import re
import sys
from scripts.trans import main
from casepunc.utils import WordTokenizer
if __name__ == '__main__':
    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
    sys.exit(main())

记得给脚本加可执行权限。

• 修改setup.py，删除entry_points里的console_scripts配置，改用scripts参数指定自定义脚本：

setup(
    # 其余配置不变
    scripts=["bin/your-cli-name"]
)

这个方案的缺点是不会自动生成Windows平台下的exe启动包装，跨平台兼容性差，仅适合纯类Unix环境使用。

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