fss-mini-rag-github/mini_rag/venv_checker.py

#!/usr/bin/env python3
"""
Virtual Environment Checker
Ensures scripts run in proper Python virtual environment for consistency and safety.
"""

import sys
import os
import sysconfig
from pathlib import Path

def is_in_virtualenv() -> bool:
    """Check if we're running in a virtual environment."""
    # Check for virtual environment indicators
    return (
        hasattr(sys, 'real_prefix') or  # virtualenv
        (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix) or  # venv/pyvenv
        os.environ.get('VIRTUAL_ENV') is not None  # Environment variable
    )

def get_expected_venv_path() -> Path:
    """Get the expected virtual environment path for this project."""
    # Assume .venv in the same directory as the script
    script_dir = Path(__file__).parent.parent
    return script_dir / '.venv'

def check_correct_venv() -> tuple[bool, str]:
    """
    Check if we're in the correct virtual environment.
    
    Returns:
        (is_correct, message)
    """
    if not is_in_virtualenv():
        return False, "not in virtual environment"
    
    expected_venv = get_expected_venv_path()
    if not expected_venv.exists():
        return False, "expected virtual environment not found"
    
    current_venv = os.environ.get('VIRTUAL_ENV')
    if current_venv:
        current_venv_path = Path(current_venv).resolve()
        expected_venv_path = expected_venv.resolve()
        
        if current_venv_path != expected_venv_path:
            return False, f"wrong virtual environment (using {current_venv_path}, expected {expected_venv_path})"
    
    return True, "correct virtual environment"

def show_venv_warning(script_name: str = "script") -> None:
    """Show virtual environment warning with helpful instructions."""
    expected_venv = get_expected_venv_path()
    
    print("⚠️  VIRTUAL ENVIRONMENT WARNING")
    print("=" * 50)
    print()
    print(f"This {script_name} should be run in a Python virtual environment for:")
    print("  • Consistent dependencies")
    print("  • Isolated package versions") 
    print("  • Proper security isolation")
    print("  • Reliable functionality")
    print()
    
    if expected_venv.exists():
        print("✅ Virtual environment found!")
        print(f"   Location: {expected_venv}")
        print()
        print("🚀 To activate it:")
        print(f"   source {expected_venv}/bin/activate")
        print(f"   {script_name}")
        print()
        print("🔄 Or run with activation:")
        print(f"   source {expected_venv}/bin/activate && {script_name}")
    else:
        print("❌ No virtual environment found!")
        print()
        print("🛠️  Create one first:")
        print("   ./install_mini_rag.sh")
        print()
        print("📚 Or manually:")
        print(f"   python3 -m venv {expected_venv}")
        print(f"   source {expected_venv}/bin/activate")
        print("   pip install -r requirements.txt")
    
    print()
    print("💡 Why this matters:")
    print("   Without a virtual environment, you may experience:")
    print("   • Import errors from missing packages")
    print("   • Version conflicts with system Python")
    print("   • Inconsistent behavior across systems")
    print("   • Potential system-wide package pollution")
    print()

def check_and_warn_venv(script_name: str = "script", force_exit: bool = False) -> bool:
    """
    Check virtual environment and warn if needed.
    
    Args:
        script_name: Name of the script for user-friendly messages
        force_exit: Whether to exit if not in correct venv
        
    Returns:
        True if in correct venv, False otherwise
    """
    is_correct, message = check_correct_venv()
    
    if not is_correct:
        show_venv_warning(script_name)
        
        if force_exit:
            print(f"⛔ Exiting {script_name} for your safety.")
            print("   Please activate the virtual environment and try again.")
            sys.exit(1)
        else:
            print(f"⚠️  Continuing anyway, but {script_name} may not work correctly...")
            print()
            return False
    
    return True

def require_venv(script_name: str = "script") -> None:
    """Require virtual environment or exit."""
    check_and_warn_venv(script_name, force_exit=True)

# Quick test function
def main():
    """Test the virtual environment checker."""
    print("🧪 Virtual Environment Checker Test")
    print("=" * 40)
    
    print(f"In virtual environment: {is_in_virtualenv()}")
    print(f"Expected venv path: {get_expected_venv_path()}")
    
    is_correct, message = check_correct_venv()
    print(f"Correct venv: {is_correct} ({message})")
    
    if not is_correct:
        show_venv_warning("test script")

if __name__ == "__main__":
    main()