python自动使用虚拟环境和安装依赖

代码如下，Windows环境测试通过（1. 判断是否在虚拟环境里；2. 判断当前目录下是否有venv文件夹；3. 如果都没有则通过python -m venv venv来创建；4. 然后调用venv里的pip来安装模块）

import platform
import re
import subprocess
import sys
from contextlib import AbstractContextManager
from pathlib import Path


class EnsureImport(AbstractContextManager):
    """Auto install modules if import error.

    Usage::
        >>> for _ range(EnsureImport.retry):
        ...     with EnsureImport(
        ...         multipart='python-multipart', dotenv='python-dotenv'
        ...     ) as _m:
        ...         import six
        ...         import multipart
        ...         from dotenv import load_dotenv
        ...         # more imports ...
        ...     if _m.ok:
        ...         break
        ...
    """

    mapping = {
        "multipart": "python-multipart",
        "dotenv": "python-dotenv",
        "snap7": "python-snap7",
    }
    retry = 30

    def __init__(self, **kwargs):
        self.exception = None
        self._success = True
        self.package_mapping = dict(self.mapping, **kwargs)

    @property
    def ok(self) -> bool:
        return self._success

    def __exit__(self, exc_type, exc_value, traceback):
        if isinstance(exc_value, (ImportError, ModuleNotFoundError)):
            self.exception = exc_value
            self._success = False
            self.run()
            return True

    def run(self):
        e = self.exception
        modules = re.findall(r"'([a-zA-Z][0-9a-zA-Z_-]+)'", str(e))
        if "--no-install" in sys.argv or not modules:
            raise e
        ms = (self.package_mapping.get(i, i) for i in modules)
        rc = self.install_and_extend_sys_path(*ms)
        if rc:
            sys.exit(rc)

    @staticmethod
    def is_venv() -> bool:
        """Whether in a virtual environment(also work for poetry)"""
        return hasattr(sys, "real_prefix") or (
            hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
        )

    @staticmethod
    def run_and_echo(cmd: str) -> int:
        print("-->\n", cmd, flush=True)
        return subprocess.call(cmd, shell=True)

    @staticmethod
    def log_error(action: str) -> None:
        print(f"ERROR: failed to {action}")

    @classmethod
    def install_and_extend_sys_path(cls, *packages) -> int:
        py = Path(sys.executable)
        depends = " ".join(packages)
        if not cls.is_venv():
            p = Path.cwd() / "venv"
            if not p.exists():
                if cls.run_and_echo(f"{py} -m venv venv"):
                    cls.log_error(f"create virtual environment for {py}")
                    return 1
            if platform.platform().lower().startswith("win"):
                py = p / "Scripts" / "python.exe"
            else:
                py = p / "bin/python"
            cls.run_and_echo(f"{py} -m pip install --upgrade pip")
            lib = list(p.rglob("site-packages"))[0]
            sys.path.append(lib.as_posix())
        if cls.run_and_echo(f"{py} -m pip install {depends}"):
            cls.log_error(f"install {depends}")
            return 2
        return 0


for _ in range(EnsureImport.retry):
    with EnsureImport(dotenv='python-dotenv') as _f:
        from dotenv import load_dotenv
    if _f.ok:
        break


load_dotenv()

 

 

#######################

以下为旧版代码，可以不用看

#######################

import os
import platform
import re
import sys
from pathlib import Path


def is_venv() -> bool:
    """判断是否处于虚拟环境(也适用于poetry的)"""
    if hasattr(sys, "real_prefix"):
        return True
    return hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix


def run_and_echo(cmd) -> int:
    print("-->\n", cmd, flush=True)
    return os.system(cmd)


def install_and_rerun(*packages):
    py = Path(sys.executable)
    if not is_venv():
        if not (p := Path("venv")).exists():
            if run_and_echo(f"{py} -m venv venv"):
                return 1
        if platform.platform().lower().startswith("win"):
            py = p / "Scripts" / "python.exe"
        else:
            py = p / "bin/python"
    if run_and_echo(f"{py} -m pip install {' '.join(packages)}"):
        return 2
    cmd = f"{py} {sys.argv[0]} --no-install {' '.join(sys.argv[1:])}"
    return run_and_echo(cmd)


try:
    import kivy
except ImportError as e:
    if "--no-install" in sys.argv:
        raise e
    modules = re.findall(r"'([a-zA-Z_-]+)'", str(e))
    sys.exit(install_and_rerun(*modules))


def main():
    pass

if __name__ == '__main__':
    main()　

运行时，如果import失败，会判断是否处于虚拟环境，是的话，直接pip install报错的缺失包，然后自动重新执行脚本;

否则，判断当前路径是否有venv文件，有的话使用venv/*/python，否则使用python -m venv venv创建它

优化成class以便PyCharm可以折叠

import os
import platform
import random
import re
import sys
from pathlib import Path


class EnsureImport:
    def __init__(self, e):
        self.exception = e

    def run(self):
        e = self.exception
        if "--no-install" in sys.argv:
            raise e
        modules = re.findall(r"'([a-zA-Z_-]+)'", str(e))
        sys.exit(self.install_and_rerun(*modules))

    @staticmethod
    def is_venv() -> bool:
        """判断是否处于虚拟环境(也适用于poetry的)"""
        if hasattr(sys, "real_prefix"):
            return True
        return hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix

    @staticmethod
    def run_and_echo(cmd) -> int:
        print("-->\n", cmd, flush=True)
        return os.system(cmd)

    @classmethod
    def install_and_rerun(cls, *packages):
        py = Path(sys.executable)
        command = " ".join(sys.argv)
        if not cls.is_venv():
            if not (p := Path("venv")).exists():
                if cls.run_and_echo(f"{py} -m venv venv"):
                    return 1
                else:
                    cls.run_and_echo(f'{py} -m pip install -U pip')
            if platform.platform().lower().startswith("win"):
                py = p / "Scripts" / "python.exe"
            else:
                py = p / "bin/python"
            return cls.run_and_echo(f"{py} {command}")
        depends = " ".join(packages)
        if cls.run_and_echo(f"{py} -m pip install {depends}"):
            return 2
        cmd = f"{py} {command} --no-install"
        return cls.run_and_echo(cmd)


try:
    import dearpygui.dearpygui as dpg
except ImportError as e:
    EnsureImport(e).run()


def main():
    print(dpg)


if __name__ == "__main__":
    main()

　　使用了isort+black+ruff进行代码格式化

进一步优化成with风格：

import os
import platform
import re
import sys
from pathlib import Path


class EnsureImport:
    """Auto install module if import error.

    Usage::
        >>> with EnsureImport():
        ...     import six

    """

    def __init__(self):
        self.exception = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            if isinstance(exc_value, (ImportError, ModuleNotFoundError)):
                self.exception = exc_value
                self.run()

    def run(self):
        e = self.exception
        if "--no-install" in sys.argv:
            raise e
        modules = re.findall(r"'([a-zA-Z_-]+)'", str(e))
        sys.exit(self.install_and_rerun(*modules))

    @staticmethod
    def is_venv() -> bool:
        """判断是否处于虚拟环境(也适用于poetry的)"""
        if hasattr(sys, "real_prefix"):
            return True
        return hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix

    @staticmethod
    def run_and_echo(cmd) -> int:
        print("-->\n", cmd, flush=True)
        return os.system(cmd)

    @classmethod
    def install_and_rerun(cls, *packages):
        py = Path(sys.executable)
        command = " ".join(sys.argv)
        if not cls.is_venv():
            if not (p := Path("venv")).exists():
                if cls.run_and_echo(f"{py} -m venv venv"):
                    return 1
            if platform.platform().lower().startswith("win"):
                py = p / "Scripts" / "python.exe"
            else:
                py = p / "bin/python"
            return cls.run_and_echo(f"{py} {command}")
        depends = " ".join(packages)
        if cls.run_and_echo(f"{py} -m pip install {depends}"):
            return 2
        cmd = f"{py} {command} --no-install"
        return cls.run_and_echo(cmd)


with EnsureImport():
    import six


def main():
    print(six.__file__)


if __name__ == "__main__":
    main()

　　进一步优化成支持多import且只run一次的：

import platform
import re
import subprocess
import sys
from contextlib import AbstractContextManager
from pathlib import Path


class EnsureImport(AbstractContextManager):
    """Auto install modules if import error.

    Usage::
        >>> for _ range(EnsureImport.retry):
        ...     with EnsureImport(
        ...         multipart='python-multipart', dotenv='python-dotenv'
        ...     ) as _m:
        ...         import six
        ...         import multipart
        ...         from dotenv import load_dotenv
        ...         # more imports ...
        ...     if _m.ok:
        ...         break
        ...
    """

    retry = 10

    def __init__(self, **kwargs):
        self.exception = None
        self._success = True
        self.package_mapping = kwargs

    @property
    def ok(self) -> bool:
        return self._success

    def __exit__(self, exc_type, exc_value, traceback):
        if isinstance(exc_value, (ImportError, ModuleNotFoundError)):
            self.exception = exc_value
            self._success = False
            self.run()
            return True

    def run(self):
        e = self.exception
        if "--no-install" in sys.argv:
            raise e
        modules = re.findall(r"'([a-zA-Z_-]+)'", str(e))
        if mp := self.package_mapping:
            modules = [mp.get(i, i) for i in modules]
        if rc := self.install_and_extend_sys_path(*modules):
            sys.exit(rc)

    @staticmethod
    def is_venv() -> bool:
        """判断是否处于虚拟环境(也适用于poetry的)"""
        if hasattr(sys, "real_prefix"):
            return True
        return hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix

    @staticmethod
    def run_and_echo(cmd: str) -> int:
        print("-->\n", cmd, flush=True)
        return subprocess.call(cmd, shell=True)

    @staticmethod
    def log_error(action: str) -> None:
        print(f"ERROR: failed to {action}")

    @classmethod
    def install_and_extend_sys_path(cls, *packages) -> int:
        py = Path(sys.executable)
        depends = " ".join(packages)
        if not cls.is_venv():
            if not (p := Path("venv")).exists():
                if depends.lower() == 'dearpygui' and sys.version_info >= (3, 11):
                    py = Path('python3.10')
                if cls.run_and_echo(f"{py} -m venv venv"):
                    cls.log_error(f"create virtual environment for {py}")
                    return 1
            if platform.platform().lower().startswith("win"):
                py = p / "Scripts" / "python.exe"
            else:
                py = p / "bin/python"
            lib = list(p.rglob("site-packages"))[0]
            sys.path.append(lib.as_posix())
        if cls.run_and_echo(f"{py} -m pip install {depends}"):
            cls.log_error(f"install {depends}")
            return 2
        return 0


for _ in range(EnsureImport.retry):
    with EnsureImport(multipart="python-multipart", dotenv="python-dotenv") as _m:
        import multipart
        import six
        from dotenv import load_dotenv
    if _m.ok:
        break


def main():
    print(six.__file__)
    print(multipart.__file__)
    load_dotenv()


if __name__ == "__main__":
    main()