"""Environment variable management for MCP Python REPL server."""
import logging
import os
import sys
import importlib
import inspect
from pathlib import Path
from typing import Dict, List, Optional, Any
from dotenv import load_dotenv
from .models import EnvironmentVariable
logger = logging.getLogger(__name__)
class EnvHandler:
"""Handles environment variable operations."""
def __init__(self, project_dir: Path | None = None):
"""Initialize EnvHandler with project directory.
Args:
project_dir: Project directory path. If None, uses current working directory.
"""
self.project_dir = (project_dir or Path.cwd()).resolve()
self.session_env_vars: Dict[str, str] = {} # Session-specific env vars
async def load_env_file(self, env_file_path: str = ".env") -> Dict[str, Any]:
"""Load environment variables from .env file.
Args:
env_file_path: Path to .env file relative to project directory
Returns:
Dictionary with load result
"""
try:
env_path = self.project_dir / env_file_path
if not env_path.exists():
return {
"success": False,
"message": f"Environment file '{env_file_path}' not found",
"path": str(env_path),
"loaded_count": 0
}
# Change to project directory before loading
original_cwd = os.getcwd()
try:
os.chdir(self.project_dir)
# Load the .env file
loaded = load_dotenv(env_path, override=True)
if loaded:
# Count loaded variables by comparing before/after
# This is approximate since load_dotenv doesn't return count
env_vars = {}
with open(env_path, 'r') as f:
for line in f:
line = line.strip()
if line and not line.startswith('#') and '=' in line:
key, _ = line.split('=', 1)
env_vars[key.strip()] = True
return {
"success": True,
"message": f"Successfully loaded environment variables from '{env_file_path}'",
"path": str(env_path),
"loaded_count": len(env_vars),
"variables": list(env_vars.keys())
}
else:
return {
"success": False,
"message": f"Failed to load '{env_file_path}' (file may be empty or invalid)",
"path": str(env_path),
"loaded_count": 0
}
finally:
os.chdir(original_cwd)
except Exception as e:
logger.error(f"Environment file load error: {e}")
return {
"success": False,
"message": f"Error loading environment file: {str(e)}",
"error": str(e),
"loaded_count": 0
}
async def set_env_var(self, name: str, value: str, session_specific: bool = True) -> Dict[str, Any]:
"""Set an environment variable.
Args:
name: Variable name
value: Variable value
session_specific: Whether to set as session-specific or global
Returns:
Dictionary with operation result
"""
try:
if session_specific:
# Store in session-specific vars
self.session_env_vars[name] = value
# Also set in os.environ for immediate use
os.environ[name] = value
return {
"success": True,
"message": f"Set session-specific environment variable '{name}'",
"name": name,
"value": value,
"session_specific": True
}
else:
# Set directly in os.environ
os.environ[name] = value
return {
"success": True,
"message": f"Set global environment variable '{name}'",
"name": name,
"value": value,
"session_specific": False
}
except Exception as e:
logger.error(f"Environment variable set error: {e}")
return {
"success": False,
"message": f"Failed to set environment variable '{name}': {str(e)}",
"error": str(e)
}
async def list_env_vars(self, include_system: bool = False) -> List[EnvironmentVariable]:
"""List current environment variables.
Args:
include_system: Whether to include system environment variables
Returns:
List of EnvironmentVariable objects
"""
try:
env_vars = []
# Add session-specific variables
for name, value in self.session_env_vars.items():
env_vars.append(EnvironmentVariable(
name=name,
value=value,
session_specific=True
))
# Add other environment variables if requested
if include_system:
for name, value in os.environ.items():
# Skip if already added as session-specific
if name not in self.session_env_vars:
env_vars.append(EnvironmentVariable(
name=name,
value=value,
session_specific=False
))
else:
# Only add project-related variables (common ones)
project_vars = [
'DATABASE_URL', 'SECRET_KEY', 'DEBUG', 'ENVIRONMENT',
'API_KEY', 'JWT_SECRET', 'REDIS_URL', 'RABBITMQ_URL',
'PORT', 'HOST', 'PYTHONPATH'
]
for name in project_vars:
if name in os.environ and name not in self.session_env_vars:
env_vars.append(EnvironmentVariable(
name=name,
value=os.environ[name],
session_specific=False
))
# Sort by name for consistent output
env_vars.sort(key=lambda x: x.name.lower())
return env_vars
except Exception as e:
logger.error(f"Environment variable list error: {e}")
return []
async def create_env_template(self, settings_class_path: str, output_path: str = ".env.example") -> Dict[str, Any]:
"""Generate .env.example from pydantic Settings model.
Args:
settings_class_path: Python path to Settings class (e.g., "app.config.Settings")
output_path: Output file path relative to project directory
Returns:
Dictionary with operation result
"""
try:
# Parse the settings class path
if '.' not in settings_class_path:
raise ValueError(f"Invalid settings class path: '{settings_class_path}'. Expected format: 'module.ClassName'")
module_path, class_name = settings_class_path.rsplit(".", 1)
# Change to project directory for import
original_cwd = os.getcwd()
original_path = sys.path.copy()
try:
os.chdir(self.project_dir)
# Add project directory to Python path
if str(self.project_dir) not in sys.path:
sys.path.insert(0, str(self.project_dir))
# Import the module
module = importlib.import_module(module_path)
settings_class = getattr(module, class_name)
# Check if it's a pydantic Settings class
# We'll check for common pydantic Settings characteristics
if not hasattr(settings_class, '__fields__') and not hasattr(settings_class, 'model_fields'):
raise ValueError(f"{settings_class_path} does not appear to be a pydantic model")
# Generate template content
lines = [
"# Auto-generated environment variable template",
f"# Generated from {settings_class_path}",
f"# Edit this file and copy to .env",
""
]
# Get fields - handle both pydantic v1 and v2
fields = {}
if hasattr(settings_class, '__fields__'):
# Pydantic v1
fields = settings_class.__fields__
elif hasattr(settings_class, 'model_fields'):
# Pydantic v2
fields = settings_class.model_fields
# Generate field entries
for field_name, field_info in fields.items():
# Get field type and description
if hasattr(field_info, 'annotation'):
# Pydantic v2
field_type = getattr(field_info, 'annotation', 'Unknown')
description = getattr(field_info, 'description', None)
else:
# Pydantic v1
field_type = getattr(field_info, 'type_', 'Unknown')
description = getattr(field_info.field_info, 'description', None) if hasattr(field_info, 'field_info') else None
# Format type name
type_name = getattr(field_type, '__name__', str(field_type))
# Add field to template
lines.append(f"# {field_name.upper()}")
if description:
lines.append(f"# Description: {description}")
lines.append(f"# Type: {type_name}")
lines.append(f"{field_name.upper()}=")
lines.append("")
# Write template file
template_path = self.project_dir / output_path
template_content = "\n".join(lines)
template_path.write_text(template_content, encoding='utf-8')
return {
"success": True,
"message": f"Generated environment template at '{output_path}'",
"path": str(template_path),
"settings_class": settings_class_path,
"field_count": len(fields),
"fields": list(fields.keys())
}
finally:
os.chdir(original_cwd)
sys.path[:] = original_path
except ImportError as e:
return {
"success": False,
"message": f"Failed to import settings class '{settings_class_path}': {str(e)}",
"error": f"ImportError: {str(e)}"
}
except Exception as e:
logger.error(f"Environment template creation error: {e}")
return {
"success": False,
"message": f"Failed to create environment template: {str(e)}",
"error": str(e)
}
def cleanup_session_vars(self):
"""Clean up session-specific environment variables."""
try:
# Remove session-specific vars from os.environ
for name in self.session_env_vars:
if name in os.environ:
del os.environ[name]
# Clear session vars
self.session_env_vars.clear()
logger.info("Cleaned up session-specific environment variables")
except Exception as e:
logger.error(f"Environment cleanup error: {e}")
def get_project_env_info(self) -> Dict[str, Any]:
"""Get information about project environment files.
Returns:
Dictionary with environment file information
"""
try:
env_files = ['.env', '.env.local', '.env.example', '.env.template']
file_info = {}
for env_file in env_files:
file_path = self.project_dir / env_file
file_info[env_file] = {
"exists": file_path.exists(),
"path": str(file_path),
"size": file_path.stat().st_size if file_path.exists() else 0
}
return {
"project_dir": str(self.project_dir),
"env_files": file_info,
"session_vars_count": len(self.session_env_vars),
"session_vars": list(self.session_env_vars.keys())
}
except Exception as e:
logger.error(f"Environment info error: {e}")
return {
"project_dir": str(self.project_dir),
"error": str(e)
}
