From 44eb15874858d8436011aaccf68eb3f3769f032b Mon Sep 17 00:00:00 2001 From: Nick Brassel Date: Mon, 17 Jun 2024 16:47:08 +1000 Subject: [PATCH] Mucking about with forcing the use of a venv at runtime. --- qmk_cli/script_qmk.py | 2 + qmk_cli/subcommands/__init__.py | 1 + qmk_cli/subcommands/pip.py | 72 +++++++++++++++++++++++++++++++++ qmk_cli/switch_to_venv.py | 37 +++++++++++++++++ 4 files changed, 112 insertions(+) create mode 100644 qmk_cli/subcommands/pip.py create mode 100644 qmk_cli/switch_to_venv.py diff --git a/qmk_cli/script_qmk.py b/qmk_cli/script_qmk.py index 7d68f76..460809f 100644 --- a/qmk_cli/script_qmk.py +++ b/qmk_cli/script_qmk.py @@ -10,6 +10,8 @@ import sys from platform import platform from traceback import print_exc +from . import switch_to_venv # noqa, intentionally unused + import milc from . import __version__ diff --git a/qmk_cli/subcommands/__init__.py b/qmk_cli/subcommands/__init__.py index 9ba4579..2252513 100644 --- a/qmk_cli/subcommands/__init__.py +++ b/qmk_cli/subcommands/__init__.py @@ -6,3 +6,4 @@ from . import clone # noqa from . import console # noqa from . import env # noqa from . import setup # noqa +from . import pip # noqa diff --git a/qmk_cli/subcommands/pip.py b/qmk_cli/subcommands/pip.py new file mode 100644 index 0000000..e88cc02 --- /dev/null +++ b/qmk_cli/subcommands/pip.py @@ -0,0 +1,72 @@ +"""Installs python dependencies. +""" +import shlex +import shutil +import subprocess +from milc import cli + +from qmk_cli.helpers import find_qmk_firmware, is_qmk_firmware + + +def _run_pip(pip_args, dry_run=False, upgrade=False): + pip_exe = ['pip'] + if cli.config.user.pip_exe: + pip_exe = shlex.split(cli.config.user.pip_exe) + pip_command = [*pip_exe, *pip_args] + pip_command[0] = shutil.which(pip_command[0]) + if shutil.which(pip_command[0]) is None: + pip_command[0] = 'pip' # Hope falling back to unqualified path works + + if dry_run: + print(" ".join(shlex.quote(str(s)) for s in pip_command)) + return True + + try: + with subprocess.Popen(pip_command, stderr=subprocess.STDOUT, stdout=subprocess.PIPE, bufsize=1, universal_newlines=True, encoding='utf-8') as p: + for line in p.stdout: + print(line, end='') + except Exception as e: + pip_cmd = ' '.join([s.replace(' ', r'\ ') for s in pip_command]) + cli.log.error("Could not run '%s': %s: %s", pip_cmd, e.__class__.__name__, e) + return False + + if p.returncode == 0: + return True + return False + + +def install_qmk_python_deps(upgrade=False, developer=False, dry_run=False): + qmk_firmware_dir = find_qmk_firmware() + if not is_qmk_firmware(qmk_firmware_dir): + cli.log.error('Could not find {fg_cyan}qmk_firmware{fg_reset}!') + return False + + reqs_path = qmk_firmware_dir / 'requirements.txt' if not developer else qmk_firmware_dir / 'requirements-dev.txt' + + pip_args = ['install', '-r', reqs_path] + if upgrade: + pip_args.append('--upgrade') + + if _run_pip(pip_args, dry_run=dry_run): + if not dry_run: + cli.log.info('Successfully installed python dependencies!') + return True + else: + if not dry_run: + cli.log.error('Failed to install python dependencies!') + return False + + +@cli.argument('-u', '--upgrade', arg_only=True, action='store_true', help='Upgrades existing packages.') +@cli.argument('-d', '--developer', arg_only=True, action='store_true', help='Installs QMK Firmware developer dependencies, too.') +@cli.argument('-n', '--dry-run', arg_only=True, action='store_true', help='Prints the commands to execute, rather than executing them.') +@cli.subcommand('Installs python dependencies.') +def install_deps(cli): + return install_qmk_python_deps(upgrade=cli.args.upgrade, developer=cli.args.developer, dry_run=cli.args.dry_run) + + +@cli.argument('-n', '--dry-run', arg_only=True, action='store_true', help='Prints the commands to execute, rather than executing them.') +@cli.argument('arguments', arg_only=True, default=[], nargs='*', help='The arguments to pass to pip.') +@cli.subcommand('Runs pip inside the qmk venv.') +def pip(cli): + return _run_pip(cli.args.arguments, dry_run=cli.args.dry_run) diff --git a/qmk_cli/switch_to_venv.py b/qmk_cli/switch_to_venv.py new file mode 100644 index 0000000..5350138 --- /dev/null +++ b/qmk_cli/switch_to_venv.py @@ -0,0 +1,37 @@ +""" Activates a virtualenv if not already present +""" + +import os +import sys + + +if not os.environ.get('VIRTUAL_ENV') or True: + import site + from pathlib import Path + + # Build a venv if required + INTENDED_VENV_PATH = Path(os.environ['HOME']) / ".local/qmk_venv" + if not INTENDED_VENV_PATH.exists(): + from venv import EnvBuilder + builder = EnvBuilder(with_pip=True) + builder.create(str(INTENDED_VENV_PATH)) + + # Determine the venv's bin directory + bin_dir = INTENDED_VENV_PATH / "bin" + + # Update the environment to point to the QMK venv (effectively matches `source path_to_venv/bin/activate`) + os.environ["PATH"] = os.pathsep.join([str(bin_dir), *os.environ.get("PATH", "").split(os.pathsep)]) + os.environ["VIRTUAL_ENV"] = str(INTENDED_VENV_PATH) + if 'PYTHONHOME' in os.environ: + os.environ.pop('PYTHONHOME') + + # Prepend the venv's library paths to the python import mechanism + pyver = sys.version_info[:2] + lib_dir = INTENDED_VENV_PATH / f'lib/python{pyver[0]}.{pyver[1]}/site-packages' + for d in [lib_dir, bin_dir]: + site.addsitedir(str(d)) + sys.path.insert(0, str(d)) + + # Update the python prefixes + sys.base_prefix = sys.prefix + sys.prefix = str(INTENDED_VENV_PATH)