venv_manager.py

"""Virtual Environment Manager for Python Scripts.

This module provides automatic virtual environment management for Python scripts,
with cross-platform support for Linux, macOS, and Windows. It handles creation,
updating, and execution of scripts within isolated virtual environments.

Features:
- Automatic venv creation and management
- Cross-platform compatibility (Linux, macOS, Windows)
- Version-pinned dependencies
- Cache-based venv storage
- Automatic requirements tracking
- Automatic cleanup of unused environments (after 30 days of inactivity)
- Platform-specific paths:
  * Linux: ~/.cache/python-venv/
  * macOS: ~/Library/Caches/python-venv/
  * Windows: %LOCALAPPDATA%/python-venv/

Example:
    Basic usage in your script (recommended):

    >>> from venv_manager import VenvManager
    >>>
    >>> REQUIREMENTS = [
    ...     "requests==2.28.1",
    ...     "pyyaml==6.0",
    ... ]
    >>>
    >>> def main():
    ...     # Your code here
    ...     pass
    >>>
    >>> if __name__ == "__main__":
    ...     VenvManager.bootstrap(REQUIREMENTS)
    ...     main()

    Advanced usage with explicit control:

    >>> venv = VenvManager(REQUIREMENTS)
    >>> if not VenvManager.is_venv():
    ...     venv.ensure_venv()
    ...     venv.restart_in_venv()

Features Details:
- Virtual environments are stored in platform-specific cache directories
- Environments are uniquely identified by script path hash
- Requirements are tracked and environments are recreated when they change
- Unused environments are automatically cleaned up after 30 days
- Weekly cleanup process removes old environments
- Cross-platform support with appropriate paths and commands
- Silent operation on Windows (no console popups)

Maintainer:
    Piotr Rogowski <[email protected]>

Source:
    https://gist.github.com/Szpadel/43794d606d9924e7fea3e63fb80042bf
"""

import os
import sys
import venv
import hashlib
import subprocess
import platform
import time
from datetime import datetime, timedelta
from pathlib import Path

# Windows-specific imports
if platform.system() == 'Windows':
    import ctypes
    from subprocess import CREATE_NO_WINDOW
else:
    CREATE_NO_WINDOW = 0  # Dummy value for non-Windows platforms

def parse_requirement(req: str) -> tuple[str, str | None]:
    """Parse requirement string into (package, version)"""
    if "==" in req:
        package, version = req.split("==", 1)
        return package.lower(), version
    return req.lower(), None

class VenvManager:
    """Manages virtual environments for Python scripts.

    This class handles the creation, validation, and execution of Python scripts
    within isolated virtual environments. It automatically manages dependencies
    and ensures consistent execution environments across different platforms.

    Args:
        requirements (list[str]): List of pip requirements, optionally version-pinned
                                (e.g., ["package==1.0.0", "other-package"])

    Attributes:
        venv_path (Path): Location of the virtual environment
        cache_dir (Path): Base directory for storing virtual environments
        requirements (list[str]): List of required packages

    Example:
        >>> venv = VenvManager(["colorama==0.4.6"])
        >>> venv.ensure_venv()  # Creates or updates venv
        >>> VenvManager.is_venv()  # Checks if running in venv
        True

    The manager includes automatic cleanup of unused virtual environments:
    - Environments unused for 30 days are automatically removed
    - Cleanup runs at most once per week
    - Usage is tracked by monitoring file access times
    - Current environment is marked as used during bootstrap
    """
    # Class-level prefix for all output messages
    _MSG_PREFIX = "[venv] "
    _CLEANUP_AFTER_DAYS = 30
    _CLEANUP_INTERVAL_DAYS = 7

    @classmethod
    def _print(cls, message: str) -> None:
        """Print a message with the venv manager prefix."""
        print(cls._MSG_PREFIX + message)

    def __init__(self, requirements: list[str]):
        self.requirements = requirements

        # Platform specific cache directory
        if platform.system() == 'Darwin':
            cache_base = Path.home() / 'Library' / 'Caches'
        elif platform.system() == 'Windows':
            cache_base = Path(os.getenv('LOCALAPPDATA', Path.home() / 'AppData' / 'Local'))
        else:
            cache_base = Path(os.getenv('XDG_CACHE_HOME', Path.home() / '.cache'))

        self.cache_dir = cache_base / 'python-venv'

        self.venv_path = self.cache_dir / f"venv-{self._get_script_hash()}"
        self.bin_dir = 'Scripts' if platform.system() == 'Windows' else 'bin'

        # Try cleanup on initialization
        self._try_cleanup()


    def _run_subprocess(self, cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
        """Run subprocess with platform-specific settings."""
        if platform.system() == 'Windows':
            kwargs['creationflags'] = CREATE_NO_WINDOW

        return subprocess.run(
            cmd,
            check=True,
            encoding='utf-8',
            errors='replace',  # Handle encoding errors gracefully
            **kwargs
        )

    def _get_script_hash(self):
        """Calculate hash based on script's absolute path"""
        return hashlib.sha256(str(Path(sys.argv[0]).resolve()).encode()).hexdigest()[:8]

    def _get_requirements_hash(self):
        """Calculate hash of requirements list"""
        requirements_str = ','.join(sorted(self.requirements))
        return hashlib.sha256(requirements_str.encode()).hexdigest()[:8]

    def _check_venv_requirements(self):
        """Check if venv requirements match current requirements"""
        req_hash_file = self.venv_path / ".requirements.hash"
        current_hash = self._get_requirements_hash()

        if not req_hash_file.exists():
            return False

        stored_hash = req_hash_file.read_text().strip()
        return stored_hash == current_hash

    def _get_pip_path(self) -> Path:
        """Get platform specific pip path"""
        return self.venv_path / self.bin_dir / ('pip.exe' if platform.system() == 'Windows' else 'pip')

    def _get_python_path(self) -> Path:
        """Get platform specific python path"""
        return self.venv_path / self.bin_dir / ('python.exe' if platform.system() == 'Windows' else 'python')

    def _check_installed_packages(self):
        """Check if all requirements are already installed with correct versions"""
        pip_path = self._get_pip_path()
        try:
            output = self._run_subprocess(
                [str(pip_path), "freeze"],
                capture_output=True,
                text=True
            ).stdout.splitlines()
            installed = {line.split('==')[0].lower(): line.split('==')[1] for line in output if '==' in line}

            for req in self.requirements:
                package, required_version = parse_requirement(req)
                if package not in installed:
                    return False
                if required_version and installed[package] != required_version:
                    return False
            return True
        except subprocess.CalledProcessError:
            return False

    @staticmethod
    def is_venv() -> bool:
        """Check if the current Python interpreter is running in a virtual environment.

        Returns:
            bool: True if running in a virtual environment, False otherwise
        """
        return hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)

    def ensure_venv(self) -> None:
        """Ensure virtual environment exists and has all required packages.

        Creates a new virtual environment if it doesn't exist, or updates an existing
        one if requirements have changed. Handles package installation and updates.

        Note:
            This method is usually called automatically by `bootstrap()`.
        """
        if self.venv_path.exists() and not self._check_venv_requirements():
            self._print("Requirements changed, recreating virtual environment...")
            import shutil
            shutil.rmtree(self.venv_path)

        if not self.venv_path.exists():
            self._print("Creating virtual environment...")
            venv.create(self.venv_path, with_pip=True)
            needs_install = True
        else:
            needs_install = not self._check_installed_packages()

        if needs_install:
            self._print("Installing required packages...")
            pip_path = self._get_pip_path()
            self._run_subprocess([str(pip_path), "install", "-U", "pip"] + self.requirements)

        req_hash_file = self.venv_path / ".requirements.hash"
        req_hash_file.write_text(self._get_requirements_hash())

    def restart_in_venv(self):
        """Restart the script inside the venv"""
        venv_python = self._get_python_path()
        os.execv(str(venv_python), [str(venv_python)] + sys.argv)

    def _mark_venv_used(self) -> None:
        """Mark the virtual environment as recently used."""
        usage_marker = self.venv_path / ".last_used"
        self.venv_path.mkdir(parents=True, exist_ok=True)
        usage_marker.touch()

    def _try_cleanup(self) -> None:
        """Try to clean up old virtual environments if enough time has passed."""
        cleanup_marker = self.cache_dir / ".cleanup_marker"
        now = datetime.now()

        # Check if we should run cleanup
        if cleanup_marker.exists():
            last_cleanup = datetime.fromtimestamp(cleanup_marker.stat().st_mtime)
            if now - last_cleanup < timedelta(days=self._CLEANUP_INTERVAL_DAYS):
                return

        self._print("Running cleanup of old virtual environments...")
        cleanup_count = 0

        # Ensure cache dir exists
        self.cache_dir.mkdir(parents=True, exist_ok=True)

        # Look for venv directories
        for venv_dir in self.cache_dir.glob("venv-*"):
            if not venv_dir.is_dir():
                continue

            # Check the most recent mtime among all files in the venv
            try:
                usage_marker = venv_dir / ".last_used"
                if usage_marker.exists():
                    last_used = datetime.fromtimestamp(usage_marker.stat().st_mtime)
                else:
                    # Fallback to checking all files if no marker exists
                    most_recent = max(
                        f.stat().st_mtime
                        for f in venv_dir.rglob("*")
                        if f.is_file()
                    )
                    last_used = datetime.fromtimestamp(most_recent)

                if now - last_used > timedelta(days=self._CLEANUP_AFTER_DAYS):
                    self._print(f"Removing unused environment: {venv_dir.name}")
                    import shutil
                    shutil.rmtree(venv_dir)
                    cleanup_count += 1

            except (OSError, ValueError):
                # If we can't check the timestamps or there's an error, skip this directory
                continue

        # Update cleanup marker
        cleanup_marker.touch()

        if cleanup_count > 0:
            self._print(f"Cleaned up {cleanup_count} unused environments")

    @staticmethod
    def bootstrap(requirements: list[str]) -> 'VenvManager':
        """Bootstrap a virtual environment and restart the script if needed.

        This is the recommended way to use VenvManager. It handles all the necessary
        steps to ensure your script runs in a virtual environment with the required
        packages.

        Args:
            requirements: List of pip requirements (e.g., ["package==1.0.0"])

        Returns:
            VenvManager: Instance of the manager, useful for further operations

        Example:
            >>> def main():
            ...     print("Running in venv")
            >>>
            >>> if __name__ == "__main__":
            ...     venv = VenvManager.bootstrap(["colorama==0.4.6"])
            ...     main()

        Note:
            Each time bootstrap is called, it marks the environment as recently used,
            preventing it from being cleaned up if it's actively used.
        """
        venv = VenvManager(requirements)
        if not VenvManager.is_venv():
            venv.ensure_venv()
            venv.restart_in_venv()
        venv._mark_venv_used()  # Mark venv as used even if we didn't need to recreate it
        return venv