pip_utils: utilities for managing python packages via pip
This module implements both syncronous and async calls to pip, separating it from the rest of the application. The syncronous implementation has no dependencies on Moonraker. Signed-off-by: Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Eric Callahan
parent
f1de614027
commit
20871e2171
|
|
@ -0,0 +1,247 @@
|
||||||
|
# Utilities for managing python packages using Pip
|
||||||
|
#
|
||||||
|
# Copyright (C) 2024 Eric Callahan <arksine.code@gmail.com>
|
||||||
|
#
|
||||||
|
# This file may be distributed under the terms of the GNU GPLv3 license
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import shlex
|
||||||
|
import subprocess
|
||||||
|
import pathlib
|
||||||
|
import shutil
|
||||||
|
import threading
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
# Annotation imports
|
||||||
|
from typing import (
|
||||||
|
TYPE_CHECKING,
|
||||||
|
Any,
|
||||||
|
Optional,
|
||||||
|
Union,
|
||||||
|
Dict,
|
||||||
|
List,
|
||||||
|
Tuple,
|
||||||
|
Callable,
|
||||||
|
IO
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from ..server import Server
|
||||||
|
from ..components.shell_command import ShellCommandFactory
|
||||||
|
|
||||||
|
MIN_PIP_VERSION = (23, 3, 2)
|
||||||
|
MIN_PYTHON_VERSION = (3, 7)
|
||||||
|
|
||||||
|
# Synchronous Subprocess Helpers
|
||||||
|
def _run_subprocess_with_response(
|
||||||
|
cmd: str,
|
||||||
|
timeout: Optional[float] = None,
|
||||||
|
env: Optional[Dict[str, str]] = None
|
||||||
|
) -> str:
|
||||||
|
prog = shlex.split(cmd)
|
||||||
|
proc = subprocess.run(
|
||||||
|
prog, capture_output=True, timeout=timeout, env=env,
|
||||||
|
check=True, text=True, errors="ignore", encoding="utf-8"
|
||||||
|
)
|
||||||
|
if proc.returncode == 0:
|
||||||
|
return proc.stdout.strip()
|
||||||
|
err = proc.stderr
|
||||||
|
raise Exception(f"Failed to run pip command '{cmd}': {err}")
|
||||||
|
|
||||||
|
def _process_subproc_output(
|
||||||
|
stdout: IO[str],
|
||||||
|
callback: Callable[[str], None]
|
||||||
|
) -> None:
|
||||||
|
for line in stdout:
|
||||||
|
callback(line.rstrip("\n"))
|
||||||
|
|
||||||
|
def _run_subprocess(
|
||||||
|
cmd: str,
|
||||||
|
timeout: Optional[float] = None,
|
||||||
|
env: Optional[Dict[str, str]] = None,
|
||||||
|
response_cb: Optional[Callable[[str], None]] = None
|
||||||
|
) -> None:
|
||||||
|
prog = shlex.split(cmd)
|
||||||
|
params: Dict[str, Any] = {"errors": "ignore", "encoding": "utf-8"}
|
||||||
|
if response_cb is not None:
|
||||||
|
params = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
|
||||||
|
with subprocess.Popen(prog, text=True, env=env, **params) as process:
|
||||||
|
if process.stdout is not None and response_cb is not None:
|
||||||
|
reader_thread = threading.Thread(
|
||||||
|
target=_process_subproc_output, args=(process.stdout, response_cb)
|
||||||
|
)
|
||||||
|
reader_thread.start()
|
||||||
|
reader_thread.join(timeout)
|
||||||
|
if reader_thread.is_alive():
|
||||||
|
process.kill()
|
||||||
|
elif timeout is not None:
|
||||||
|
process.wait(timeout)
|
||||||
|
ret = process.poll()
|
||||||
|
if ret != 0:
|
||||||
|
raise Exception(f"Failed to run pip command '{cmd}'")
|
||||||
|
|
||||||
|
@ dataclass(frozen=True)
|
||||||
|
class PipVersionInfo:
|
||||||
|
pip_version_string: str
|
||||||
|
python_version_string: str
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pip_version(self) -> Tuple[int, ...]:
|
||||||
|
return tuple(int(part) for part in self.pip_version_string.split("."))
|
||||||
|
|
||||||
|
@property
|
||||||
|
def python_version(self) -> Tuple[int, ...]:
|
||||||
|
return tuple(int(part) for part in self.python_version_string.split("."))
|
||||||
|
|
||||||
|
class PipExecutor:
|
||||||
|
def __init__(
|
||||||
|
self, pip_cmd: str, response_handler: Optional[Callable[[str], None]] = None
|
||||||
|
) -> None:
|
||||||
|
self.pip_cmd = pip_cmd
|
||||||
|
self.response_hdlr = response_handler
|
||||||
|
|
||||||
|
def call_pip_with_response(
|
||||||
|
self,
|
||||||
|
args: str,
|
||||||
|
timeout: Optional[float] = None,
|
||||||
|
env: Optional[Dict[str, str]] = None
|
||||||
|
) -> str:
|
||||||
|
return _run_subprocess_with_response(f"{self.pip_cmd} {args}", timeout, env)
|
||||||
|
|
||||||
|
def call_pip(
|
||||||
|
self,
|
||||||
|
args: str,
|
||||||
|
timeout: Optional[float] = None,
|
||||||
|
env: Optional[Dict[str, str]] = None
|
||||||
|
) -> None:
|
||||||
|
_run_subprocess(f"{self.pip_cmd} {args}", timeout, env, self.response_hdlr)
|
||||||
|
|
||||||
|
def get_pip_version(self) -> PipVersionInfo:
|
||||||
|
resp = self.call_pip_with_response("--version", 10.)
|
||||||
|
return parse_pip_version(resp)
|
||||||
|
|
||||||
|
def update_pip(self) -> None:
|
||||||
|
pip_ver = ".".join([str(part) for part in MIN_PIP_VERSION])
|
||||||
|
self.call_pip(f"install pip=={pip_ver}", 120.)
|
||||||
|
|
||||||
|
def install_packages(
|
||||||
|
self,
|
||||||
|
packages: Union[pathlib.Path, List[str]],
|
||||||
|
sys_env_vars: Optional[Dict[str, Any]] = None
|
||||||
|
) -> None:
|
||||||
|
args = prepare_install_args(packages)
|
||||||
|
env: Optional[Dict[str, str]] = None
|
||||||
|
if sys_env_vars is not None:
|
||||||
|
env = dict(os.environ)
|
||||||
|
env.update(sys_env_vars)
|
||||||
|
self.call_pip(f"install {args}", timeout=1200., env=env)
|
||||||
|
|
||||||
|
def build_virtualenv(self, py_exec: pathlib.Path, args: str) -> None:
|
||||||
|
bin_dir = py_exec.parent
|
||||||
|
env_path = bin_dir.parent.resolve()
|
||||||
|
if env_path.exists():
|
||||||
|
shutil.rmtree(env_path)
|
||||||
|
_run_subprocess(
|
||||||
|
f"virtualenv {args} {env_path}",
|
||||||
|
timeout=600.,
|
||||||
|
response_cb=self.response_hdlr
|
||||||
|
)
|
||||||
|
if not py_exec.exists():
|
||||||
|
raise Exception("Failed to create new virtualenv", 500)
|
||||||
|
|
||||||
|
class AsyncPipExecutor:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
pip_cmd: str,
|
||||||
|
server: Server,
|
||||||
|
notify_callback: Optional[Callable[[bytes], None]] = None
|
||||||
|
) -> None:
|
||||||
|
self.pip_cmd = pip_cmd
|
||||||
|
self.server = server
|
||||||
|
self.notify_callback = notify_callback
|
||||||
|
|
||||||
|
def get_shell_cmd(self) -> ShellCommandFactory:
|
||||||
|
return self.server.lookup_component("shell_command")
|
||||||
|
|
||||||
|
async def get_pip_version(self) -> PipVersionInfo:
|
||||||
|
resp: str = await self.get_shell_cmd().exec_cmd(
|
||||||
|
f"{self.pip_cmd} --version", timeout=30., attempts=3, log_stderr=True
|
||||||
|
)
|
||||||
|
return parse_pip_version(resp)
|
||||||
|
|
||||||
|
async def update_pip(self) -> None:
|
||||||
|
pip_ver = ".".join([str(part) for part in MIN_PIP_VERSION])
|
||||||
|
shell_cmd = self.get_shell_cmd()
|
||||||
|
await shell_cmd.run_cmd_async(
|
||||||
|
f"{self.pip_cmd} install pip=={pip_ver}",
|
||||||
|
self.notify_callback, timeout=1200., attempts=3, log_stderr=True
|
||||||
|
)
|
||||||
|
|
||||||
|
async def install_packages(
|
||||||
|
self,
|
||||||
|
packages: Union[pathlib.Path, List[str]],
|
||||||
|
sys_env_vars: Optional[Dict[str, Any]] = None
|
||||||
|
) -> None:
|
||||||
|
# Update python dependencies
|
||||||
|
args = prepare_install_args(packages)
|
||||||
|
env: Optional[Dict[str, str]] = None
|
||||||
|
if sys_env_vars is not None:
|
||||||
|
env = dict(os.environ)
|
||||||
|
env.update(sys_env_vars)
|
||||||
|
shell_cmd = self.get_shell_cmd()
|
||||||
|
await shell_cmd.run_cmd_async(
|
||||||
|
f"{self.pip_cmd} install {args}", self.notify_callback,
|
||||||
|
timeout=1200., attempts=3, env=env, log_stderr=True
|
||||||
|
)
|
||||||
|
|
||||||
|
async def build_virtualenv(self, py_exec: pathlib.Path, args: str) -> None:
|
||||||
|
bin_dir = py_exec.parent
|
||||||
|
env_path = bin_dir.parent.resolve()
|
||||||
|
if env_path.exists():
|
||||||
|
shutil.rmtree(env_path)
|
||||||
|
shell_cmd = self.get_shell_cmd()
|
||||||
|
await shell_cmd.exec_cmd(f"virtualenv {args} {env_path}", timeout=600.)
|
||||||
|
if not py_exec.exists():
|
||||||
|
raise self.server.error("Failed to create new virtualenv", 500)
|
||||||
|
|
||||||
|
def read_requirements_file(requirements_path: pathlib.Path) -> List[str]:
|
||||||
|
if not requirements_path.is_file():
|
||||||
|
raise FileNotFoundError(f"Requirements file {requirements_path} not found")
|
||||||
|
data = requirements_path.read_text()
|
||||||
|
modules: List[str] = []
|
||||||
|
for line in data.split("\n"):
|
||||||
|
line = line.strip()
|
||||||
|
if not line or line[0] in "#-":
|
||||||
|
continue
|
||||||
|
match = re.search(r"\s#", line)
|
||||||
|
if match is not None:
|
||||||
|
line = line[:match.start()].strip()
|
||||||
|
modules.append(line)
|
||||||
|
return modules
|
||||||
|
|
||||||
|
def parse_pip_version(pip_response: str) -> PipVersionInfo:
|
||||||
|
match = re.match(
|
||||||
|
r"^pip ([0-9.]+) from .+? \(python ([0-9.]+)\)$", pip_response.strip()
|
||||||
|
)
|
||||||
|
if match is None:
|
||||||
|
raise ValueError("Unable to parse pip version from response")
|
||||||
|
pipver_str: str = match.group(1).strip()
|
||||||
|
pyver_str: str = match.group(2).strip()
|
||||||
|
return PipVersionInfo(pipver_str, pyver_str)
|
||||||
|
|
||||||
|
def check_pip_needs_update(version_info: PipVersionInfo) -> bool:
|
||||||
|
if version_info.python_version < MIN_PYTHON_VERSION:
|
||||||
|
return False
|
||||||
|
return version_info.pip_version < MIN_PIP_VERSION
|
||||||
|
|
||||||
|
def prepare_install_args(packages: Union[pathlib.Path, List[str]]) -> str:
|
||||||
|
if isinstance(packages, pathlib.Path):
|
||||||
|
if not packages.is_file():
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"Invalid path to requirements_file '{packages}'"
|
||||||
|
)
|
||||||
|
return f"-r {packages}"
|
||||||
|
reqs = [req.replace("\"", "'") for req in packages]
|
||||||
|
return " ".join([f"\"{req}\"" for req in reqs])
|
||||||
Loading…
Reference in New Issue