#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Virtual environment utilities.
:note: Reference implementation.
:date: 2021
:author: Christian Wiche
:contact: cwichel@gmail.com
:license: The MIT License (MIT)
"""
# -------------------------------------
import contextlib as ctx
import os
import site
import sys
import typing as tp
import attr
from .common import TPAny
from .path import Path
# -->> Tunables <<---------------------
# -->> Definitions <<------------------
@attr.define
class VENV:
"""
Virtual environment info.
"""
os_venv: str = attr.field(default="")
sys_prefix: str = attr.field(default=None)
#: History of activated environments
VENVS: tp.List[VENV] = []
# -->> API <<--------------------------
def venv_activate(path: TPAny) -> None:
"""
Activates the given virtual environment on the current run.
:param TPAny path: Path to virtualenv folder.
:raises EnvironmentError: Virtualenv is not found, or it doesn't match target.
"""
# Path Validation
path = Path.validate_dir(path=path, must_exist=True)
if not list(path.rglob(pattern="python.exe")):
raise EnvironmentError(f"Python interpreter not found on virtual environment: {path}")
# Add current environment to env list
VENVS.append(VENV(os_venv=os.environ.get("VIRTUAL_ENV", ""), sys_prefix=sys.prefix))
# Register and include binaries
os.environ["VIRTUAL_ENV"] = f"{path}"
os.environ["PATH"] = f"{(path / 'Scripts')}{os.pathsep}{os.environ['PATH']}"
# Add python libraries to host
idx = len(sys.path)
site.addsitedir(sitedir=str(path / "Libs/site-packages"))
sys.path = sys.path[idx:] + sys.path[0:idx]
# Update system prefix
sys.prefix = f"{path}"
# Verify
active = Path(sys.prefix)
if path != active:
VENVS.pop(-1)
raise EnvironmentError(f"Active environment doesn't match target: {path} != {active}.")
def venv_deactivate() -> None:
"""
Deactivates the current virtual environment.
:note: This function can only deactivate environments activated during execution (stored in VENVS).
"""
# Deactivate (only if we activated it)
if VENVS:
old_venv = VENVS.pop(-1)
# Unregister environment
os.environ["VIRTUAL_ENV"] = old_venv.os_venv
# Remove binaries and libraries
os.environ["PATH"] = f"{os.pathsep}".join(os.environ['PATH'].split(os.pathsep)[1:])
sys.path = sys.path[1:]
# Restore system prefix
sys.prefix = old_venv.sys_prefix
[docs]@ctx.contextmanager
def env(*remove, **update) -> tp.Iterator[None]:
"""
Temporarily updates the "os.environ" dictionary in-place.
The "os.environ" dictionary is updated in-place so that the modification
is sure to work in all situations.
:param remove: Environment variables to remove.
:param update: Dictionary of environment variables and values to add/update.
"""
# Get variables
cwe = os.environ
update = update or {}
remove = remove or []
# List of environment variables being updated or removed.
stomped = (set(update.keys()) | set(remove)) & set(cwe.keys())
# Environment variables and values to restore on exit.
update_after = {k: cwe[k] for k in stomped}
# Environment variables and values to remove on exit.
remove_after = frozenset(k for k in update if k not in cwe)
# Execute
try:
# Update and execute
cwe.update(update)
_ = [cwe.pop(k, None) for k in remove]
yield
finally:
# Restore
cwe.update(update_after)
_ = [cwe.pop(k) for k in remove_after]
[docs]@ctx.contextmanager
def venv(path: TPAny) -> tp.Iterator[None]:
"""
Temporarily activates the given virtual environment.
:param TPAny path: Path to virtualenv folder.
"""
try:
venv_activate(path=path)
yield
finally:
venv_deactivate()
# -->> Export <<-----------------------
__all__ = [
"env",
"venv",
]