Spaces:
Sleeping
Sleeping
| """ | |
| 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__ | |
| 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 | |
| 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__}>" | |