File size: 8,401 Bytes
cacd4d0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
"""
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__}>"