Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
mkvenv: avoid ensurepip if pip is installed
Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
  • Loading branch information
jnsnow authored and bonzini committed May 16, 2023
1 parent 8616c06 commit f1b1892
Showing 1 changed file with 64 additions and 3 deletions.
67 changes: 64 additions & 3 deletions python/scripts/mkvenv.py
Expand Up @@ -11,6 +11,8 @@
Commands:
command Description
create create a venv
post_init
post-venv initialization
ensure Ensure that the specified package is installed.
--------------------------------------------------
Expand All @@ -25,6 +27,13 @@
--------------------------------------------------
usage: mkvenv post_init [-h]
options:
-h, --help show this help message and exit
--------------------------------------------------
usage: mkvenv ensure [-h] [--online] [--dir DIR] dep_spec...
positional arguments:
Expand Down Expand Up @@ -111,7 +120,9 @@ class QemuEnvBuilder(venv.EnvBuilder):
The primary difference is that it emulates a "nested" virtual
environment when invoked from inside of an existing virtual
environment by including packages from the parent.
environment by including packages from the parent. Also,
"ensurepip" is replaced if possible with just recreating pip's
console_scripts inside the virtual environment.
Parameters for base class init:
- system_site_packages: bool = False
Expand All @@ -138,8 +149,19 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
# The venv we are currently in, also does so.
kwargs["system_site_packages"] = sys.base_prefix in site.PREFIXES

if kwargs.get("with_pip", False):
check_ensurepip()
# ensurepip is slow: venv creation can be very fast for cases where
# we allow the use of system_site_packages. Therefore, ensurepip is
# replaced with our own script generation once the virtual environment
# is setup.
self.want_pip = kwargs.get("with_pip", False)
if self.want_pip:
if (
kwargs.get("system_site_packages", False)
and not need_ensurepip()
):
kwargs["with_pip"] = False
else:
check_ensurepip()

super().__init__(*args, **kwargs)

Expand Down Expand Up @@ -211,6 +233,14 @@ def post_post_setup(self, context: SimpleNamespace) -> None:
with open(pth_file, "w", encoding="UTF-8") as file:
file.write(parent_libpath + os.linesep)

if self.want_pip:
args = [
context.env_exe,
__file__,
"post_init",
]
subprocess.run(args, check=True)

def get_value(self, field: str) -> str:
"""
Get a string value from the context namespace after a call to build.
Expand All @@ -223,6 +253,19 @@ def get_value(self, field: str) -> str:
return ret


def need_ensurepip() -> bool:
"""
Tests for the presence of setuptools and pip.
:return: `True` if we do not detect both packages.
"""
# Don't try to actually import them, it's fraught with danger:
# https://github.com/pypa/setuptools/issues/2993
if find_spec("setuptools") and find_spec("pip"):
return False
return True


def check_ensurepip() -> None:
"""
Check that we have ensurepip.
Expand Down Expand Up @@ -693,6 +736,17 @@ def ensure(
raise SystemExit(f"\n{msg}\n\n") from exc


def post_venv_setup() -> None:
"""
This is intended to be run *inside the venv* after it is created.
"""
logger.debug("post_venv_setup()")
# Generate a 'pip' script so the venv is usable in a normal
# way from the CLI. This only happens when we inherited pip from a
# parent/system-site and haven't run ensurepip in some way.
generate_console_scripts(["pip"])


def _add_create_subcommand(subparsers: Any) -> None:
subparser = subparsers.add_parser("create", help="create a venv")
subparser.add_argument(
Expand All @@ -703,6 +757,10 @@ def _add_create_subcommand(subparsers: Any) -> None:
)


def _add_post_init_subcommand(subparsers: Any) -> None:
subparsers.add_parser("post_init", help="post-venv initialization")


def _add_ensure_subcommand(subparsers: Any) -> None:
subparser = subparsers.add_parser(
"ensure", help="Ensure that the specified package is installed."
Expand Down Expand Up @@ -761,6 +819,7 @@ def main() -> int:
)

_add_create_subcommand(subparsers)
_add_post_init_subcommand(subparsers)
_add_ensure_subcommand(subparsers)

args = parser.parse_args()
Expand All @@ -771,6 +830,8 @@ def main() -> int:
system_site_packages=True,
clear=True,
)
if args.command == "post_init":
post_venv_setup()
if args.command == "ensure":
ensure(
dep_specs=args.dep_specs,
Expand Down

0 comments on commit f1b1892

Please sign in to comment.