mermaid-trace 0.4.1__py3-none-any.whl → 0.5.3.post0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,104 +1,275 @@
1
1
  """
2
2
  Function Tracing Decorator Module
3
-
4
- This module provides the core tracing functionality for MermaidTrace. It includes:
5
- - A decorator (`trace`/`trace_interaction`) to instrument functions for tracing
6
- - Helper functions for formatting and logging trace events
7
- - Context management for tracking call chains
8
-
9
- The decorator can be applied to both synchronous and asynchronous functions,
10
- and automatically handles context propagation for nested calls.
3
+ =================================
4
+
5
+ This module provides the core tracing functionality for MermaidTrace. It implements the decorators
6
+ responsible for intercepting function calls, capturing execution details, and logging them as
7
+ structured events that can be visualized as Mermaid Sequence Diagrams.
8
+
9
+ Key Components:
10
+ ---------------
11
+ 1. **`@trace` Decorator**: The primary interface for users. It can be used as a simple decorator
12
+ `@trace` or with arguments `@trace(action="Login")`.
13
+ 2. **Context Management**: Handles the propagation of "who called whom". It uses `LogContext`
14
+ to track the current "participant" (caller) so that nested function calls correctly
15
+ link to their parent caller.
16
+ 3. **Event Logging**: Generates `FlowEvent` objects containing source, target, action, and
17
+ parameters, which are then passed to the specialized logging handler.
18
+ 4. **Automatic Target Resolution**: Heuristics to intelligently guess the participant name
19
+ from the class instance (`self`) or class (`cls`), falling back to the module name.
20
+
21
+ Usage:
22
+ ------
23
+ from mermaid_trace.core.decorators import trace
24
+
25
+ @trace
26
+ def my_function(x):
27
+ return x + 1
28
+
29
+ @trace(target="Database", action="Query Users")
30
+ async def get_users():
31
+ ...
11
32
  """
12
33
 
13
34
  import functools
14
35
  import logging
15
36
  import inspect
37
+ import re
16
38
  import reprlib
17
- from typing import Optional, Any, Callable, Tuple, Dict, Union, TypeVar, cast, overload
39
+ import traceback
40
+ from dataclasses import dataclass
41
+ from typing import (
42
+ Optional,
43
+ Any,
44
+ Callable,
45
+ Tuple,
46
+ Dict,
47
+ Union,
48
+ TypeVar,
49
+ cast,
50
+ overload,
51
+ List,
52
+ )
18
53
 
19
54
  from .events import FlowEvent
20
55
  from .context import LogContext
56
+ from .config import config
21
57
 
22
- # Logger name for flow events - used to isolate tracing logs from other application logs
58
+ # Logger name for flow events - used to isolate tracing logs from other application logs.
59
+ # This specific name is often used to configure a separate file handler in logging configs.
23
60
  FLOW_LOGGER_NAME = "mermaid_trace.flow"
24
61
 
25
- # Define generic type variable for the decorated function
62
+ # Define generic type variable for the decorated function to preserve type hints
26
63
  F = TypeVar("F", bound=Callable[..., Any])
27
64
 
28
65
 
29
66
  def get_flow_logger() -> logging.Logger:
30
- """Returns the dedicated logger for flow events.
67
+ """
68
+ Returns the dedicated logger for flow events.
69
+
70
+ This logger is intended to be used only for emitting `FlowEvent` objects.
71
+ It separates tracing noise from standard application logging.
31
72
 
32
73
  Returns:
33
- logging.Logger: Logger instance configured for tracing events
74
+ logging.Logger: Logger instance configured for tracing events.
34
75
  """
35
76
  return logging.getLogger(FLOW_LOGGER_NAME)
36
77
 
37
78
 
38
- def _safe_repr(obj: Any, max_len: int = 50, max_depth: int = 1) -> str:
79
+ class FlowRepr(reprlib.Repr):
80
+ """
81
+ Custom Repr class that simplifies default Python object representations
82
+ (those containing memory addresses) into a cleaner <ClassName> format.
83
+ Also groups consecutive identical items in lists to keep diagrams concise.
39
84
  """
40
- Safely creates a string representation of an object.
41
85
 
42
- Prevents massive log files by truncating long strings/objects
43
- and handling exceptions during __repr__ calls (e.g. strict objects).
86
+ def _group_items(self, items_str: List[str]) -> List[str]:
87
+ """Groups consecutive identical strings in a list."""
88
+ if not items_str:
89
+ return []
90
+ res = []
91
+ current_item = items_str[0]
92
+ current_count = 1
93
+ for i in range(1, len(items_str)):
94
+ if items_str[i] == current_item:
95
+ current_count += 1
96
+ else:
97
+ if current_count > 1:
98
+ res.append(f"{current_item} x {current_count}")
99
+ else:
100
+ res.append(current_item)
101
+ current_item = items_str[i]
102
+ current_count = 1
103
+ # Handle the last group
104
+ if current_count > 1:
105
+ res.append(f"{current_item} x {current_count}")
106
+ else:
107
+ res.append(current_item)
108
+ return res
109
+
110
+ def repr_list(self, obj: List[Any], level: int) -> str:
111
+ """Custom list representation with item grouping."""
112
+ n = len(obj)
113
+ if n == 0:
114
+ return "[]"
115
+ items_str = []
116
+ for i in range(min(n, self.maxlist)):
117
+ items_str.append(self.repr1(obj[i], level - 1))
118
+
119
+ grouped = self._group_items(items_str)
120
+ if n > self.maxlist:
121
+ grouped.append("...")
122
+ return "[" + ", ".join(grouped) + "]"
123
+
124
+ def repr_tuple(self, obj: Tuple[Any, ...], level: int) -> str:
125
+ """Custom tuple representation with item grouping."""
126
+ n = len(obj)
127
+ if n == 0:
128
+ return "()"
129
+ if n == 1:
130
+ return "(" + self.repr1(obj[0], level - 1) + ",)"
131
+ items_str = []
132
+ for i in range(min(n, self.maxtuple)):
133
+ items_str.append(self.repr1(obj[i], level - 1))
134
+
135
+ grouped = self._group_items(items_str)
136
+ if n > self.maxtuple:
137
+ grouped.append("...")
138
+ return "(" + ", ".join(grouped) + ")"
139
+
140
+ def repr1(self, x: Any, level: int) -> str:
141
+ # Check if the object uses the default object.__repr__
142
+ # Default repr looks like <module.Class object at 0x...>
143
+ raw = repr(x)
144
+ if " object at 0x" in raw and raw.startswith("<") and raw.endswith(">"):
145
+ # Simplify to just <ClassName>
146
+ return f"<{x.__class__.__name__}>"
147
+
148
+ # Also handle some common cases where address is present but it's not the default repr
149
+ if " at 0x" in raw and raw.startswith("<") and raw.endswith(">"):
150
+ # Try to extract the class name or just simplify it
151
+ return f"<{x.__class__.__name__}>"
152
+
153
+ return super().repr1(x, level)
154
+
155
+
156
+ def _safe_repr(
157
+ obj: Any, max_len: Optional[int] = None, max_depth: Optional[int] = None
158
+ ) -> str:
159
+ """
160
+ Safely creates a string representation of an object for logging purposes.
161
+
162
+ This function is critical for preventing log bloat and runtime errors during tracing.
163
+ It handles:
164
+ 1. **Truncation**: Limits the length of strings to prevent huge log files.
165
+ 2. **Depth Control**: Limits recursion for nested structures like dicts/lists.
166
+ 3. **Error Handling**: Catches exceptions if an object's `__repr__` is buggy or strict.
167
+ 4. **Simplification**: Automatically simplifies default Python object representations
168
+ (containing memory addresses) into <ClassName>.
44
169
 
45
170
  Args:
46
- obj: The object to represent as a string
47
- max_len: Maximum length of the resulting string
48
- max_depth: Maximum recursion depth for nested objects
171
+ obj: The object to represent as a string.
172
+ max_len: Maximum length of the resulting string before truncation.
173
+ Defaults to config.max_string_length if None.
174
+ max_depth: Maximum recursion depth for nested objects.
175
+ Defaults to config.max_arg_depth if None.
49
176
 
50
177
  Returns:
51
- str: Safe, truncated representation of the object
178
+ str: Safe, truncated representation of the object.
52
179
  """
53
- try:
54
- # Create a custom repr object to control depth and length
55
- a_repr = reprlib.Repr()
56
- a_repr.maxstring = max_len
57
- a_repr.maxother = max_len
58
- a_repr.maxlevel = max_depth
180
+ # Use config defaults if not explicitly provided
181
+ final_max_len = max_len if max_len is not None else config.max_string_length
182
+ final_max_depth = max_depth if max_depth is not None else config.max_arg_depth
59
183
 
184
+ try:
185
+ # Use our custom FlowRepr to provide standard way to limit representation size
186
+ # and simplify default object reprs recursively.
187
+ a_repr = FlowRepr()
188
+ a_repr.maxstring = final_max_len
189
+ a_repr.maxother = final_max_len
190
+ a_repr.maxlevel = final_max_depth
191
+
192
+ # Generate the representation
60
193
  r = a_repr.repr(obj)
61
- if len(r) > max_len:
62
- return r[:max_len] + "..."
194
+
195
+ # Final pass: Catch any remaining memory addresses using regex
196
+ # (e.g., in types reprlib doesn't recurse into)
197
+ # 1. <__main__.Class object at 0x...> -> <Class>
198
+ r = re.sub(
199
+ r"<([a-zA-Z0-9_.]+\.)?([a-zA-Z0-9_]+) object at 0x[0-9a-fA-F]+>",
200
+ r"<\2>",
201
+ r,
202
+ )
203
+ # 2. <Class at 0x...> -> <Class>
204
+ r = re.sub(
205
+ r"<([a-zA-Z0-9_.]+\.)?([a-zA-Z0-9_]+) at 0x[0-9a-fA-F]+>",
206
+ r"<\2>",
207
+ r,
208
+ )
209
+
210
+ # Double-check length constraint as reprlib might sometimes exceed it slightly
211
+ if len(r) > final_max_len:
212
+ return r[:final_max_len] + "..."
63
213
  return r
64
214
  except Exception:
65
- # Fallback if repr() fails for any reason
215
+ # Fallback if repr() fails (e.g., property access raising error in __repr__)
66
216
  return "<unrepresentable>"
67
217
 
68
218
 
219
+ @dataclass
220
+ class _TraceConfig:
221
+ """Internal container for tracing configuration to avoid PLR0913."""
222
+
223
+ capture_args: Optional[bool] = None
224
+ max_arg_length: Optional[int] = None
225
+ max_arg_depth: Optional[int] = None
226
+
227
+
69
228
  def _format_args(
70
229
  args: Tuple[Any, ...],
71
230
  kwargs: Dict[str, Any],
72
- capture_args: bool = True,
73
- max_arg_length: int = 50,
74
- max_arg_depth: int = 1,
231
+ config_obj: _TraceConfig,
75
232
  ) -> str:
76
233
  """
77
- Formats function arguments into a single string "arg1, arg2, k=v".
78
- Used for the arrow label in the diagram.
234
+ Formats function arguments into a single string for the diagram arrow label.
235
+
236
+ Example Output: "1, 'test', debug=True"
237
+
238
+ This string is what appears on the arrow in the Mermaid diagram (e.g., `User->System: login(args...)`).
79
239
 
80
240
  Args:
81
- args: Positional arguments to format
82
- kwargs: Keyword arguments to format
83
- capture_args: Whether to capture and format arguments at all
84
- max_arg_length: Maximum length of each argument representation
85
- max_arg_depth: Maximum recursion depth for nested arguments
241
+ args: Positional arguments tuple.
242
+ kwargs: Keyword arguments dictionary.
243
+ config_obj: Trace configuration object.
86
244
 
87
245
  Returns:
88
- str: Formatted string of arguments, or empty string if capture_args is False
246
+ str: Comma-separated string of formatted arguments.
89
247
  """
90
- if not capture_args:
248
+ final_capture = (
249
+ config_obj.capture_args
250
+ if config_obj.capture_args is not None
251
+ else config.capture_args
252
+ )
253
+ if not final_capture:
91
254
  return ""
92
255
 
93
256
  parts: list[str] = []
94
257
 
95
- # Format positional arguments
258
+ # Process positional arguments
96
259
  for arg in args:
97
- parts.append(_safe_repr(arg, max_len=max_arg_length, max_depth=max_arg_depth))
260
+ parts.append(
261
+ _safe_repr(
262
+ arg,
263
+ max_len=config_obj.max_arg_length,
264
+ max_depth=config_obj.max_arg_depth,
265
+ )
266
+ )
98
267
 
99
- # Format keyword arguments
268
+ # Process keyword arguments
100
269
  for k, v in kwargs.items():
101
- val_str = _safe_repr(v, max_len=max_arg_length, max_depth=max_arg_depth)
270
+ val_str = _safe_repr(
271
+ v, max_len=config_obj.max_arg_length, max_depth=config_obj.max_arg_depth
272
+ )
102
273
  parts.append(f"{k}={val_str}")
103
274
 
104
275
  return ", ".join(parts)
@@ -108,76 +279,92 @@ def _resolve_target(
108
279
  func: Callable[..., Any], args: Tuple[Any, ...], target_override: Optional[str]
109
280
  ) -> str:
110
281
  """
111
- Determines the name of the participant (Target) for the diagram.
282
+ Determines the name of the 'Target' participant (the callee) for the diagram.
112
283
 
113
- Resolution Priority:
114
- 1. **Override**: If the user explicitly provided `target="Name"`, use it.
115
- 2. **Instance Method**: If the first arg looks like `self` (has __class__),
116
- use the class name.
117
- 3. **Class Method**: If the first arg is a type (cls), use the class name.
118
- 4. **Module Function**: Fallback to the name of the module containing the function.
119
- 5. **Fallback**: "Unknown".
284
+ The 'Target' is the entity *receiving* the call. We try to infer a meaningful name
285
+ (like the class name) so the diagram shows "User -> AuthService" instead of "User -> login".
286
+
287
+ Resolution Logic:
288
+ 1. **Override**: Use explicit `target` from decorator if provided.
289
+ 2. **Instance Method**: If first arg looks like `self` (has `__class__`), use ClassName.
290
+ 3. **Class Method**: If first arg is a type (cls), use ClassName.
291
+ 4. **Module Function**: Use the module name (e.g., "utils" from "my.pkg.utils").
292
+ 5. **Fallback**: "Unknown".
120
293
 
121
294
  Args:
122
- func: The function being called
123
- args: Positional arguments passed to the function
124
- target_override: Explicit target name provided by the user, if any
295
+ func: The function being called (for module inspection).
296
+ args: Positional arguments (to check for self/cls).
297
+ target_override: Explicit target name provided by user via decorator.
125
298
 
126
299
  Returns:
127
- str: Resolved target name for the diagram
300
+ str: The resolved name for the target participant.
128
301
  """
129
302
  if target_override:
130
303
  return target_override
131
304
 
132
- # Heuristic: If it's a method call, args[0] is usually 'self' or 'cls'
305
+ # Heuristic: Check if this is a method call where args[0] is 'self' or 'cls'
133
306
  if args:
134
307
  first_arg = args[0]
135
- # Check if it looks like a class instance (not a primitive or container)
308
+
309
+ # Check for class method (cls) - where first arg is the type itself
310
+ if isinstance(first_arg, type):
311
+ return first_arg.__name__
312
+
313
+ # Check for instance method (self)
314
+ # We filter out primitives because functions might take an int/str as first arg,
315
+ # which shouldn't be treated as 'self'.
136
316
  if hasattr(first_arg, "__class__") and not isinstance(
137
- first_arg, (str, int, float, bool, list, dict, type)
317
+ first_arg, (str, int, float, bool, list, dict, set, tuple)
138
318
  ):
139
319
  return str(first_arg.__class__.__name__)
140
- # Check if it looks like a class (cls) - e.g. @classmethod
141
- if isinstance(first_arg, type):
142
- return first_arg.__name__
143
320
 
144
- # Fallback to module name for standalone functions
321
+ # Fallback: Use module name for standalone functions
145
322
  module = inspect.getmodule(func)
146
323
  if module:
324
+ # Extract just the last part of the module path (e.g. 'auth' from 'app.core.auth')
147
325
  return module.__name__.split(".")[-1]
326
+
148
327
  return "Unknown"
149
328
 
150
329
 
330
+ @dataclass
331
+ class _TraceMetadata:
332
+ """Internal container for trace metadata to avoid PLR0913."""
333
+
334
+ source: str
335
+ target: str
336
+ action: str
337
+ trace_id: str
338
+
339
+
151
340
  def _log_interaction(
152
341
  logger: logging.Logger,
153
- source: str,
154
- target: str,
155
- action: str,
342
+ meta: _TraceMetadata,
156
343
  params: str,
157
- trace_id: str,
158
344
  ) -> None:
159
345
  """
160
- Logs the 'Call' event (Start of function).
161
- Generates a FlowEvent and logs it with the appropriate context.
346
+ Logs the 'Call' event (Start of function execution).
347
+
348
+ This corresponds to the solid arrow in Mermaid: `Source -> Target: Action(params)`
162
349
 
163
350
  Args:
164
- logger: Logger instance to use for logging
165
- source: Name of the source participant (caller)
166
- target: Name of the target participant (callee)
167
- action: Name of the action being performed
168
- params: Formatted string of function arguments
169
- trace_id: Unique trace identifier for correlation
351
+ logger: Logger instance.
352
+ meta: Trace metadata.
353
+ params: Stringified arguments.
170
354
  """
171
355
  req_event = FlowEvent(
172
- source=source,
173
- target=target,
174
- action=action,
175
- message=action,
356
+ source=meta.source,
357
+ target=meta.target,
358
+ action=meta.action,
359
+ message=meta.action,
176
360
  params=params,
177
- trace_id=trace_id,
361
+ trace_id=meta.trace_id,
362
+ )
363
+ # The 'extra' dict is crucial. The custom LogHandler extracts 'flow_event'
364
+ # from here to format the actual Mermaid syntax line.
365
+ logger.info(
366
+ f"{meta.source}->{meta.target}: {meta.action}", extra={"flow_event": req_event}
178
367
  )
179
- # The 'extra' dict is critical: the Handler will pick this up to format the Mermaid line
180
- logger.info(f"{source}->{target}: {action}", extra={"flow_event": req_event})
181
368
 
182
369
 
183
370
  def _log_return(
@@ -187,35 +374,43 @@ def _log_return(
187
374
  action: str,
188
375
  result: Any,
189
376
  trace_id: str,
190
- capture_args: bool = True,
191
- max_arg_length: int = 50,
192
- max_arg_depth: int = 1,
377
+ config_obj: _TraceConfig,
193
378
  ) -> None:
194
379
  """
195
- Logs the 'Return' event (End of function).
196
- Arrow: target --> source (Dotted line return)
380
+ Logs the 'Return' event (End of function execution).
197
381
 
198
- Note: 'source' here is the original caller, 'target' is the callee.
199
- So the return arrow goes from target back to source.
382
+ This corresponds to the dotted return arrow in Mermaid: `Target --> Source: Return value`
383
+
384
+ Note on Direction:
385
+ - In the diagram, the return goes from `target` (callee) back to `source` (caller).
386
+ - The code logs it as `Target->Source` to reflect this flow.
200
387
 
201
388
  Args:
202
- logger: Logger instance to use for logging
203
- source: Name of the original caller
204
- target: Name of the callee that's returning
205
- action: Name of the action being returned from
206
- result: Return value of the function
207
- trace_id: Unique trace identifier for correlation
208
- capture_args: Whether to include the return value in the log
209
- max_arg_length: Maximum length of the return value representation
210
- max_arg_depth: Maximum recursion depth for nested return values
389
+ logger: Logger instance.
390
+ source: The original caller (who will receive the return).
391
+ target: The callee (who is returning).
392
+ action: The action that is completing.
393
+ result: The return value of the function.
394
+ trace_id: Trace correlation ID.
395
+ config_obj: Trace configuration object.
211
396
  """
212
397
  result_str = ""
213
- if capture_args:
214
- result_str = _safe_repr(result, max_len=max_arg_length, max_depth=max_arg_depth)
398
+ final_capture = (
399
+ config_obj.capture_args
400
+ if config_obj.capture_args is not None
401
+ else config.capture_args
402
+ )
403
+
404
+ if final_capture:
405
+ result_str = _safe_repr(
406
+ result,
407
+ max_len=config_obj.max_arg_length,
408
+ max_depth=config_obj.max_arg_depth,
409
+ )
215
410
 
216
411
  resp_event = FlowEvent(
217
- source=target,
218
- target=source,
412
+ source=target, # Return flows FROM target
413
+ target=source, # Return flows TO source
219
414
  action=action,
220
415
  message="Return",
221
416
  is_return=True,
@@ -227,41 +422,46 @@ def _log_return(
227
422
 
228
423
  def _log_error(
229
424
  logger: logging.Logger,
230
- source: str,
231
- target: str,
232
- action: str,
425
+ meta: _TraceMetadata,
233
426
  error: Exception,
234
- trace_id: str,
235
427
  ) -> None:
236
428
  """
237
429
  Logs an 'Error' event if the function raises an exception.
238
- Arrow: target -x source (Error return)
430
+
431
+ This corresponds to the 'X' arrow in Mermaid: `Target -x Source: Error Message`
239
432
 
240
433
  Args:
241
- logger: Logger instance to use for logging
242
- source: Name of the original caller
243
- target: Name of the callee that encountered an error
244
- action: Name of the action that failed
245
- error: Exception that was raised
246
- trace_id: Unique trace identifier for correlation
434
+ logger: Logger instance.
435
+ meta: Trace metadata.
436
+ error: The exception object.
247
437
  """
438
+ # Capture full stack trace
439
+ stack_trace = "".join(
440
+ traceback.format_exception(type(error), error, error.__traceback__)
441
+ )
442
+
248
443
  err_event = FlowEvent(
249
- source=target,
250
- target=source,
251
- action=action,
444
+ source=meta.target,
445
+ target=meta.source,
446
+ action=meta.action,
252
447
  message=str(error),
253
448
  is_return=True,
254
- is_error=True,
449
+ is_error=True, # Flags this as an error event
255
450
  error_message=str(error),
256
- trace_id=trace_id,
451
+ stack_trace=stack_trace,
452
+ trace_id=meta.trace_id,
453
+ )
454
+ logger.error(
455
+ f"{meta.target}-x{meta.source}: Error", extra={"flow_event": err_event}
257
456
  )
258
- logger.error(f"{target}-x{source}: Error", extra={"flow_event": err_event})
259
457
 
260
458
 
459
+ # Overload 1: Simple usage -> @trace
261
460
  @overload
262
461
  def trace_interaction(func: F) -> F: ...
263
462
 
264
463
 
464
+ # Overload 2: Configured usage -> @trace(action="Login")
265
465
  @overload
266
466
  def trace_interaction(
267
467
  *,
@@ -269,9 +469,9 @@ def trace_interaction(
269
469
  target: Optional[str] = None,
270
470
  name: Optional[str] = None,
271
471
  action: Optional[str] = None,
272
- capture_args: bool = True,
273
- max_arg_length: int = 50,
274
- max_arg_depth: int = 1,
472
+ capture_args: Optional[bool] = None,
473
+ max_arg_length: Optional[int] = None,
474
+ max_arg_depth: Optional[int] = None,
275
475
  ) -> Callable[[F], F]: ...
276
476
 
277
477
 
@@ -282,10 +482,10 @@ def trace_interaction(
282
482
  target: Optional[str] = None,
283
483
  name: Optional[str] = None,
284
484
  action: Optional[str] = None,
285
- capture_args: bool = True,
286
- max_arg_length: int = 50,
287
- max_arg_depth: int = 1,
288
- ) -> Union[F, Callable[[F], F]]:
485
+ capture_args: Optional[bool] = None,
486
+ max_arg_length: Optional[int] = None,
487
+ max_arg_depth: Optional[int] = None,
488
+ ) -> Union[F, Callable[[F], F]]: # noqa: PLR0913
289
489
  """
290
490
  Main Decorator for tracing function execution in Mermaid diagrams.
291
491
 
@@ -294,42 +494,46 @@ def trace_interaction(
294
494
  and automatically handles context propagation for nested calls.
295
495
 
296
496
  It supports two modes of operation:
297
- 1. **Simple**: `@trace` (No arguments) - uses default settings
298
- 2. **Configured**: `@trace(action="Login", target="AuthService")` - customizes behavior
497
+ 1. **Simple Mode**: `@trace` (No arguments). Uses default naming (Class/Module name) and behavior.
498
+ 2. **Configured Mode**: `@trace(action="Login", target="AuthService")`. Customizes the diagram labels.
299
499
 
300
500
  Args:
301
- func: The function being decorated (automatically passed in simple mode).
302
- source: Explicit name of the caller participant (rarely used, usually inferred from Context).
303
- target: Explicit name of the callee participant (overrides automatic resolution).
304
- name: Alias for 'target' (for clearer API usage).
305
- action: Label for the arrow (defaults to function name in Title Case).
306
- capture_args: Whether to include arguments and return values in the log. Default True.
307
- max_arg_length: Maximum string length for argument/result representation. Default 50.
308
- max_arg_depth: Maximum recursion depth for argument/result representation. Default 1.
501
+ func: The function being decorated (passed automatically in Simple Mode).
502
+ source: Explicit name of the caller participant. Rarely used, as it's usually inferred from Context.
503
+ target: Explicit name of the callee participant. Overrides automatic `self`/`cls`/module resolution.
504
+ name: Alias for 'target' (syntactic sugar).
505
+ action: Label for the arrow. Defaults to the function name (Title Cased).
506
+ capture_args: If True, logs arguments and return values. Set False for sensitive data.
507
+ max_arg_length: Truncation limit for logging arguments/results.
508
+ max_arg_depth: Recursion limit for logging arguments/results.
309
509
 
310
510
  Returns:
311
- Callable: Either the decorated function (simple mode) or a decorator factory (configured mode)
511
+ Callable: The decorated function (in Simple Mode) or a decorator factory (in Configured Mode).
312
512
  """
313
513
 
314
- # Handle alias - 'name' is an alternative name for 'target'
514
+ # Handle alias - 'name' is an alternative convenience name for 'target'
315
515
  final_target = target or name
316
516
 
317
- # Mode 1: @trace used without parentheses - directly decorate the function
517
+ # Mode 1: @trace used without parentheses
518
+ # func is passed directly. We create the wrapper immediately.
318
519
  if func is not None and callable(func):
319
520
  return _create_decorator(
320
521
  func,
321
522
  source,
322
523
  final_target,
323
524
  action,
324
- capture_args,
325
- max_arg_length,
326
- max_arg_depth,
525
+ _TraceConfig(capture_args, max_arg_length, max_arg_depth),
327
526
  )
328
527
 
329
- # Mode 2: @trace(...) used with arguments -> returns a factory that will decorate the function
528
+ # Mode 2: @trace(...) used with arguments
529
+ # func is None. We return a "factory" function that Python will call with the function later.
330
530
  def factory(f: F) -> F:
331
531
  return _create_decorator(
332
- f, source, final_target, action, capture_args, max_arg_length, max_arg_depth
532
+ f,
533
+ source,
534
+ final_target,
535
+ action,
536
+ _TraceConfig(capture_args, max_arg_length, max_arg_depth),
333
537
  )
334
538
 
335
539
  return factory
@@ -340,64 +544,65 @@ def _create_decorator(
340
544
  source: Optional[str],
341
545
  target: Optional[str],
342
546
  action: Optional[str],
343
- capture_args: bool,
344
- max_arg_length: int,
345
- max_arg_depth: int,
547
+ config_obj: _TraceConfig,
346
548
  ) -> F:
347
549
  """
348
- Constructs the actual wrapper function for the decorated function.
349
- Handles both synchronous and asynchronous functions by creating the appropriate wrapper.
550
+ Internal factory that constructs the actual wrapper function.
350
551
 
351
- This function separates the wrapper creation logic from the argument parsing logic
352
- in `trace_interaction`, making the code cleaner and easier to test.
552
+ This separates the wrapper creation logic from the argument parsing logic in `trace_interaction`.
553
+ It handles the distinction between synchronous and asynchronous functions, returning
554
+ the appropriate wrapper type.
353
555
 
354
556
  Args:
355
- func: The function to decorate
356
- source: Explicit source name, if any
357
- target: Explicit target name, if any
358
- action: Explicit action name, if any
359
- capture_args: Whether to capture arguments and return values
360
- max_arg_length: Maximum length for argument/return value representations
361
- max_arg_depth: Maximum recursion depth for nested objects
557
+ func: The function to decorate.
558
+ source: Configured source name.
559
+ target: Configured target name.
560
+ action: Configured action name.
561
+ config_obj: Trace configuration object.
362
562
 
363
563
  Returns:
364
- Callable: Decorated function with tracing logic
564
+ Callable: The wrapped function containing tracing logic.
365
565
  """
366
566
 
367
- # Pre-calculate static metadata to save time at runtime
567
+ # Pre-calculate static metadata to save time at runtime.
568
+ # If no action name provided, generate one from the function name (e.g., "get_user" -> "Get User")
368
569
  if action is None:
369
- # Default action name is the function name, converted to Title Case
370
570
  action = func.__name__.replace("_", " ").title()
371
571
 
372
572
  @functools.wraps(func)
373
573
  def wrapper(*args: Any, **kwargs: Any) -> Any:
374
- """Synchronous function wrapper that adds tracing logic."""
574
+ """
575
+ Synchronous function wrapper.
576
+ Executes tracing logic around a standard blocking function call.
577
+ """
375
578
  # 1. Resolve Context
376
- # 'source' is who called us (from Context). 'target' is who we are (resolved from self/cls).
579
+ # 'current_source' is who called us. If not explicit, we get it from thread-local storage.
377
580
  current_source = source or LogContext.current_participant()
378
581
  trace_id = LogContext.current_trace_id()
582
+
583
+ # 'current_target' is who we are. We figure this out from 'self', 'cls', or module name.
379
584
  current_target = _resolve_target(func, args, target)
380
585
 
586
+ meta = _TraceMetadata(current_source, current_target, action, trace_id)
587
+
381
588
  logger = get_flow_logger()
382
589
  # Format arguments for the diagram arrow label
383
- params_str = _format_args(
384
- args, kwargs, capture_args, max_arg_length, max_arg_depth
385
- )
590
+ params_str = _format_args(args, kwargs, config_obj)
386
591
 
387
592
  # 2. Log Request (Start of function)
388
- # Logs the initial "Call" arrow (Source -> Target)
389
- _log_interaction(
390
- logger, current_source, current_target, action, params_str, trace_id
391
- )
593
+ # Emits the "Call" arrow (Source -> Target)
594
+ _log_interaction(logger, meta, params_str)
392
595
 
393
596
  # 3. Execute with New Context
394
- # We push 'current_target' as the NEW 'participant' (source) for any internal calls.
395
- # This builds the chain: A -> B, then inside B, B becomes the source for C (B -> C).
597
+ # We push 'current_target' as the NEW 'participant' (source) for any internal calls made by this function.
598
+ # This builds the chain: A calls B (A->B), then B calls C (B->C).
396
599
  with LogContext.scope({"participant": current_target, "trace_id": trace_id}):
397
600
  try:
601
+ # Execute the actual user function
398
602
  result = func(*args, **kwargs)
603
+
399
604
  # 4. Log Success Return
400
- # Logs the "Return" arrow (Target --> Source)
605
+ # Emits the "Return" arrow (Target --> Source)
401
606
  _log_return(
402
607
  logger,
403
608
  current_source,
@@ -405,43 +610,45 @@ def _create_decorator(
405
610
  action,
406
611
  result,
407
612
  trace_id,
408
- capture_args,
409
- max_arg_length,
410
- max_arg_depth,
613
+ config_obj,
411
614
  )
412
615
  return result
413
616
  except Exception as e:
414
617
  # 5. Log Error Return
415
- # Logs the "Error" arrow (Target --x Source)
416
- _log_error(logger, current_source, current_target, action, e, trace_id)
618
+ # Emits the "Error" arrow (Target -x Source)
619
+ _log_error(logger, meta, e)
620
+ # Re-raise the exception so program flow isn't altered
417
621
  raise
418
622
 
419
623
  @functools.wraps(func)
420
624
  async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
421
- """Asynchronous function wrapper that adds tracing logic."""
625
+ """
626
+ Asynchronous function wrapper.
627
+ Executes tracing logic around an async/await coroutine.
628
+ """
422
629
  # 1. Resolve Context (Same as sync)
423
630
  current_source = source or LogContext.current_participant()
424
631
  trace_id = LogContext.current_trace_id()
425
632
  current_target = _resolve_target(func, args, target)
426
633
 
634
+ meta = _TraceMetadata(current_source, current_target, action, trace_id)
635
+
427
636
  logger = get_flow_logger()
428
- params_str = _format_args(
429
- args, kwargs, capture_args, max_arg_length, max_arg_depth
430
- )
637
+ params_str = _format_args(args, kwargs, config_obj)
431
638
 
432
639
  # 2. Log Request
433
- _log_interaction(
434
- logger, current_source, current_target, action, params_str, trace_id
435
- )
640
+ _log_interaction(logger, meta, params_str)
436
641
 
437
642
  # 3. Execute with New Context using 'ascope'
438
- # Use async context manager (ascope) to ensure context propagates correctly across awaits.
439
- # This is critical for asyncio: context must be preserved even if the task yields control.
643
+ # Crucial difference for Async: We use `ascope` (async scope) which uses contextvars.
644
+ # This ensures the context is preserved across `await` points where the event loop might switch tasks.
440
645
  async with LogContext.ascope(
441
646
  {"participant": current_target, "trace_id": trace_id}
442
647
  ):
443
648
  try:
649
+ # Await the actual user coroutine
444
650
  result = await func(*args, **kwargs)
651
+
445
652
  # 4. Log Success Return
446
653
  _log_return(
447
654
  logger,
@@ -450,20 +657,18 @@ def _create_decorator(
450
657
  action,
451
658
  result,
452
659
  trace_id,
453
- capture_args,
454
- max_arg_length,
455
- max_arg_depth,
660
+ config_obj,
456
661
  )
457
662
  return result
458
663
  except Exception as e:
459
664
  # 5. Log Error Return
460
- _log_error(logger, current_source, current_target, action, e, trace_id)
665
+ _log_error(logger, meta, e)
461
666
  raise
462
667
 
463
- # Detect if the wrapped function is a coroutine to choose the right wrapper
668
+ # Detect if the wrapped function is a coroutine (async def)
464
669
  if inspect.iscoroutinefunction(func):
465
- return cast(F, async_wrapper) # Return async wrapper for coroutines
466
- return cast(F, wrapper) # Return sync wrapper for regular functions
670
+ return cast(F, async_wrapper) # Use async wrapper for async functions
671
+ return cast(F, wrapper) # Use sync wrapper for regular functions
467
672
 
468
673
 
469
674
  # Alias for easy import - 'trace' is the primary name users should use