Python venv injection


对我来说，这是很常见的一种情况：

  • 你的一个库，要运行一些用户的 python 的代码文件。

对我个人来说，我在 3 种情况中遇到了这个问题：

  • m 这个库，支持注册 plugin，plugin 就是在当前目录或任何sys.path中的目录创建 m/commands 文件夹，在里面的文件夹暴露一个 app。这里面的代码会被暴露为一个 entry point
  • hmr 这个库，会运行用户的文件，其衍生品uvicorn-hmr会执行用户的某个 entry point
  • embed 这个库，会执行用户 pull 下来的 app，这个 app 可能有单独的虚拟环境

最简单的解决办法就是读进来exec对吧，当然你有点经验的话就知道这样对 traceback 不友好。比较正确的办法是用runpy.run_path

但是还会遇到一个问题：

  • 用户如果这个文件用到一些包，而这些包不在你安装你这个库的环境中，导致在你的库的 python 进程中会 import 失败

最好的办法当然是在用户侧解决，比如在用户的项目中安装你这个库，但是

  1. 可能这是一个开发工具，用户不太愿意往项目中引入开发依赖
  2. 终归还是得让用户多做一步

IPython 的解决方案

如果你用过 ipython，你会发现如果你在一个没有安装 ipython 的虚拟环境中运行 ipython（你在外部安装了 ipython），它照样能 import 你虚拟环境中的包，但是会有一行 warning：

UserWarning: Attempting to work in a virtualenv. If you encounter problems, please install IPython inside the virtualenv.

去翻源码，能看到是在 这里 警告的。会出现什么 problems 呢，如果你 import 一些非纯 python 库，确实有可能出问题。复现的方式，就是你安装 ipython 的 python minor version 和虚拟环境的 python 不一样。

源码 也给出了答案，让你能 import 虚拟环境的库是靠把虚拟环境的 site-packages 也添加到当前 python 进程的 sys.path 中解决的

我也在我的m中 实现了这个逻辑，来支持运行用户在项目中定义的 m 插件，相关代码如下：

if os.getenv("VIRTUAL_ENV") and (venv_python := shutil.which("python")) and not Path(venv_python).samefile(sys.executable):
    from site import addsitepackages
    from subprocess import run

    site_packages: list[str] = eval(run([venv_python, "-c", "import site; print(repr(site.getsitepackages()))"], capture_output=True, text=True, check=True).stdout)

    path.extend(site_packages)
    addsitepackages(set(site_packages))

和 IPython 有一点区别，它是把 site-packages 与 VIRTUAL_ENV 的相对路径写死了，考虑了 Windows 和 unix 的可能性。但我不喜欢这样，我是通过在 subprocess 中执行site.getsitepackages()来获得的。这样理应更鲁棒，而且我也曾经在 pdm 的源码中看到过大量类似的使用。

更好的解决方法

我不满足于实现类似 IPython 的这种 workaround。我之前在参与 embed-ai/embed 的时候就在想，其实可以在对应的虚拟环境中启动 python，然后将自己这个库加到它的虚拟环境中再启动。这样的效果应该和“在那个虚拟环境中安装这个库”是一样的（如果我们这个库和它的依赖都是纯 Python 的）。

于是我就修改了下我的代码，在这段subprocess.run中也返回一下sys.version_info，代表 python 的版本，检查如果和当前进程不符的话就用虚拟环境的 python executable 来重新执行自己。

但是这样其实启动了两次 subprocess，并不符合我追求极致性能的惯例。转念一想，其实既然都要用 subprocess 了，不如第一次也不用了，只要用户在虚拟环境中，就一定用虚拟环境的 Python

if os.getenv("VIRTUAL_ENV") and (venv_python := shutil.which("python")) and str(Path(venv_python)) != sys.executable:
    from site import getsitepackages
    from subprocess import run

    s = set(getsitepackages())
    exit(
        run(
            [
                venv_python,
                "-c",
                f"import site,sys; sys.argv={['m', *sys.argv[1:]]}; sys.path.extend({[*s, str(Path(__file__, '../' * 3).resolve())]}); site.addsitepackages({s}); from m.cli.main import app; app()",
            ]
        ).returncode
    )

源文件在这里：utils/inject.py · CNSeniorious000/m

隐患

这样的问题在于，执行你的库的环境可能不是安装你的库的环境。这意味着如果你这个库有非纯 Python 的依赖，就很可能会报错（除非恰好用户的虚拟环境中也装了你的这个依赖）。