""" Logging Context Management. Provides context managers and decorators for: - Operation tracking with timing - Contextual logging with nested contexts - Automatic exception logging """ import logging import time import functools from contextlib import contextmanager from typing import Any, Callable, Dict, Optional, TypeVar, ParamSpec P = ParamSpec('P') R = TypeVar('R') class LogContext: """ Context manager for logging operations with timing and context. Features: - Automatic start/end logging - Timing measurement - Exception capture - Nested context support Example: logger = get_logger(__name__) with LogContext(logger, "optimization", iteration=5): # ... optimization code ... logger.info("Processing sample") # Inherits context # Output: # INFO | Starting optimization | iteration=5 # INFO | Processing sample | iteration=5 # INFO | Completed optimization | iteration=5 duration_ms=1234 """ def __init__( self, logger: logging.Logger, operation: str, log_start: bool = True, log_end: bool = True, log_level: int = logging.INFO, **context_fields: Any ): """ Initialize log context. Args: logger: Logger instance to use operation: Name of the operation being performed log_start: Whether to log when entering context log_end: Whether to log when exiting context log_level: Log level for start/end messages **context_fields: Additional fields to include in all logs """ self.logger = logger self.operation = operation self.log_start = log_start self.log_end = log_end self.log_level = log_level self.context_fields = context_fields self.start_time: Optional[float] = None self.exception: Optional[Exception] = None def __enter__(self) -> "LogContext": """Enter the context, logging start if configured.""" self.start_time = time.perf_counter() if self.log_start: self.logger.log( self.log_level, f"Starting {self.operation}", extra=self.context_fields ) return self def __exit__(self, exc_type, exc_val, exc_tb) -> bool: """Exit the context, logging completion or error.""" duration_ms = (time.perf_counter() - self.start_time) * 1000 extra = { **self.context_fields, "duration_ms": round(duration_ms, 2) } if exc_type is not None: # Log exception self.exception = exc_val self.logger.error( f"Failed {self.operation}: {exc_type.__name__}: {exc_val}", extra=extra, exc_info=True ) # Don't suppress the exception return False if self.log_end: self.logger.log( self.log_level, f"Completed {self.operation}", extra=extra ) return False def log(self, level: int, message: str, **extra_fields: Any) -> None: """Log a message within this context, inheriting context fields.""" self.logger.log( level, message, extra={**self.context_fields, **extra_fields} ) def info(self, message: str, **extra_fields: Any) -> None: """Log info message within context.""" self.log(logging.INFO, message, **extra_fields) def debug(self, message: str, **extra_fields: Any) -> None: """Log debug message within context.""" self.log(logging.DEBUG, message, **extra_fields) def warning(self, message: str, **extra_fields: Any) -> None: """Log warning message within context.""" self.log(logging.WARNING, message, **extra_fields) def error(self, message: str, **extra_fields: Any) -> None: """Log error message within context.""" self.log(logging.ERROR, message, **extra_fields) def log_operation( logger: Optional[logging.Logger] = None, operation: Optional[str] = None, log_args: bool = False, log_result: bool = False, log_level: int = logging.INFO, ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Decorator for logging function execution. Automatically logs: - Function entry (with arguments if configured) - Function exit (with result if configured) - Execution duration - Exceptions Args: logger: Logger to use (defaults to logger named after module) operation: Operation name (defaults to function name) log_args: Whether to log function arguments log_result: Whether to log function result log_level: Log level for messages Example: @log_operation(log_args=True) def process_batch(batch_id: int, items: List[str]) -> int: return len(items) # Output: # INFO | Starting process_batch | batch_id=123 items=['a', 'b'] # INFO | Completed process_batch | duration_ms=45.2 result=2 """ def decorator(func: Callable[P, R]) -> Callable[P, R]: nonlocal logger, operation if logger is None: logger = logging.getLogger(func.__module__) if operation is None: operation = func.__name__ @functools.wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: start_time = time.perf_counter() # Build context fields extra: Dict[str, Any] = {} if log_args: # Include positional args (skip self for methods) arg_names = func.__code__.co_varnames[:func.__code__.co_argcount] for i, (name, value) in enumerate(zip(arg_names, args)): if name != 'self': extra[name] = _safe_repr(value) # Include keyword args for key, value in kwargs.items(): extra[key] = _safe_repr(value) logger.log(log_level, f"Starting {operation}", extra=extra) try: result = func(*args, **kwargs) duration_ms = (time.perf_counter() - start_time) * 1000 result_extra: Dict[str, Any] = {"duration_ms": round(duration_ms, 2)} if log_result: result_extra["result"] = _safe_repr(result) logger.log(log_level, f"Completed {operation}", extra=result_extra) return result except Exception as e: duration_ms = (time.perf_counter() - start_time) * 1000 logger.error( f"Failed {operation}: {type(e).__name__}: {e}", extra={"duration_ms": round(duration_ms, 2)}, exc_info=True ) raise return wrapper return decorator @contextmanager def timed_block(logger: logging.Logger, description: str, log_level: int = logging.DEBUG): """ Simple context manager for timing a block of code. Less verbose than LogContext, suitable for quick timing measurements. Example: with timed_block(logger, "data processing"): process_data() # Output: DEBUG | data processing completed in 123.45ms """ start = time.perf_counter() try: yield finally: duration_ms = (time.perf_counter() - start) * 1000 logger.log(log_level, f"{description} completed in {duration_ms:.2f}ms") def _safe_repr(value: Any, max_length: int = 100) -> str: """ Create a safe string representation of a value for logging. Truncates long strings and handles non-serializable objects. """ try: repr_str = repr(value) if len(repr_str) > max_length: return repr_str[:max_length] + "..." return repr_str except Exception: return f"<{type(value).__name__}>"