mermaid-trace 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

mermaid_trace/core/decorators.py

@@ -12,46 +12,68 @@ FLOW_LOGGER_NAME = "mermaid_trace.flow"
 # Define generic type variable for the decorated function
 F = TypeVar("F", bound=Callable[..., Any])
 
+
 def get_flow_logger() -> logging.Logger:
     """Returns the dedicated logger for flow events."""
     return logging.getLogger(FLOW_LOGGER_NAME)
 
-def _safe_repr(obj: Any, max_len: int = 50) -> str:
+
+def _safe_repr(obj: Any, max_len: int = 50, max_depth: int = 1) -> str:
     """
     Safely creates a string representation of an object.
-    
-    Prevents massive log files by truncating long strings/objects 
+
+    Prevents massive log files by truncating long strings/objects
     and handling exceptions during __repr__ calls (e.g. strict objects).
     """
     try:
-        r = reprlib.repr(obj)
+        # Create a custom repr object to control depth and length
+        a_repr = reprlib.Repr()
+        a_repr.maxstring = max_len
+        a_repr.maxother = max_len
+        a_repr.maxlevel = max_depth
+
+        r = a_repr.repr(obj)
         if len(r) > max_len:
             return r[:max_len] + "..."
         return r
     except Exception:
         return "<unrepresentable>"
 
-def _format_args(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> str:
+
+def _format_args(
+    args: Tuple[Any, ...],
+    kwargs: Dict[str, Any],
+    capture_args: bool = True,
+    max_arg_length: int = 50,
+    max_arg_depth: int = 1,
+) -> str:
     """
     Formats function arguments into a single string "arg1, arg2, k=v".
     Used for the arrow label in the diagram.
     """
+    if not capture_args:
+        return ""
+
     parts = []
     for arg in args:
-        parts.append(_safe_repr(arg))
-        
+        parts.append(_safe_repr(arg, max_len=max_arg_length, max_depth=max_arg_depth))
+
     for k, v in kwargs.items():
-        parts.append(f"{k}={_safe_repr(v)}")
-        
+        val_str = _safe_repr(v, max_len=max_arg_length, max_depth=max_arg_depth)
+        parts.append(f"{k}={val_str}")
+
     return ", ".join(parts)
 
-def _resolve_target(func: Callable[..., Any], args: Tuple[Any, ...], target_override: Optional[str]) -> str:
+
+def _resolve_target(
+    func: Callable[..., Any], args: Tuple[Any, ...], target_override: Optional[str]
+) -> str:
     """
     Determines the name of the participant (Target) for the diagram.
-    
+
     Resolution Priority:
     1. **Override**: If the user explicitly provided `target="Name"`, use it.
-    2. **Instance Method**: If the first arg looks like `self` (has __class__), 
+    2. **Instance Method**: If the first arg looks like `self` (has __class__),
        use the class name.
     3. **Class Method**: If the first arg is a type (cls), use the class name.
     4. **Module Function**: Fallback to the name of the module containing the function.
@@ -59,18 +81,20 @@ def _resolve_target(func: Callable[..., Any], args: Tuple[Any, ...], target_over
     """
     if target_override:
         return target_override
-        
+
     # Heuristic: If it's a method call, args[0] is usually 'self'.
     if args:
         first_arg = args[0]
         # Check if it looks like a class instance
         # We check hasattr(__class__) to distinguish objects from primitives/containers broadly,
         # ensuring we don't mislabel a plain list passed as first arg to a function as a "List" participant.
-        if hasattr(first_arg, "__class__") and not isinstance(first_arg, (str, int, float, bool, list, dict, type)):
-             return str(first_arg.__class__.__name__)
+        if hasattr(first_arg, "__class__") and not isinstance(
+            first_arg, (str, int, float, bool, list, dict, type)
+        ):
+            return str(first_arg.__class__.__name__)
         # Check if it looks like a class (cls) - e.g. @classmethod
         if isinstance(first_arg, type):
-             return first_arg.__name__
+            return first_arg.__name__
 
     # Fallback to module name for standalone functions
     module = inspect.getmodule(func)
@@ -78,118 +102,180 @@ def _resolve_target(func: Callable[..., Any], args: Tuple[Any, ...], target_over
         return module.__name__.split(".")[-1]
     return "Unknown"
 
-def _log_interaction(logger: logging.Logger, 
-                     source: str, 
-                     target: str, 
-                     action: str, 
-                     params: str, 
-                     trace_id: str) -> None:
+
+def _log_interaction(
+    logger: logging.Logger,
+    source: str,
+    target: str,
+    action: str,
+    params: str,
+    trace_id: str,
+) -> None:
     """
     Logs the 'Call' event (Start of function).
     Arrow: source -> target
     """
     req_event = FlowEvent(
-        source=source, target=target, 
-        action=action, message=action,
-        params=params, trace_id=trace_id
+        source=source,
+        target=target,
+        action=action,
+        message=action,
+        params=params,
+        trace_id=trace_id,
     )
     # The 'extra' dict is critical: the Handler will pick this up to format the Mermaid line
     logger.info(f"{source}->{target}: {action}", extra={"flow_event": req_event})
 
-def _log_return(logger: logging.Logger, 
-                source: str, 
-                target: str, 
-                action: str, 
-                result: Any, 
-                trace_id: str) -> None:
+
+def _log_return(
+    logger: logging.Logger,
+    source: str,
+    target: str,
+    action: str,
+    result: Any,
+    trace_id: str,
+    capture_args: bool = True,
+    max_arg_length: int = 50,
+    max_arg_depth: int = 1,
+) -> None:
     """
     Logs the 'Return' event (End of function).
     Arrow: target --> source (Dotted line return)
-    
+
     Note: 'source' here is the original caller, 'target' is the callee.
     So the return arrow goes from target back to source.
     """
-    result_str = _safe_repr(result)
+    result_str = ""
+    if capture_args:
+        result_str = _safe_repr(result, max_len=max_arg_length, max_depth=max_arg_depth)
+
     resp_event = FlowEvent(
-        source=target, target=source, 
-        action=action, message="Return", 
-        is_return=True, result=result_str, trace_id=trace_id
+        source=target,
+        target=source,
+        action=action,
+        message="Return",
+        is_return=True,
+        result=result_str,
+        trace_id=trace_id,
     )
     logger.info(f"{target}->{source}: Return", extra={"flow_event": resp_event})
 
-def _log_error(logger: logging.Logger, 
-               source: str, 
-               target: str, 
-               action: str, 
-               error: Exception, 
-               trace_id: str) -> None:
+
+def _log_error(
+    logger: logging.Logger,
+    source: str,
+    target: str,
+    action: str,
+    error: Exception,
+    trace_id: str,
+) -> None:
     """
     Logs an 'Error' event if the function raises an exception.
     Arrow: target -x source (Error return)
     """
     err_event = FlowEvent(
-        source=target, target=source, action=action, 
-        message=str(error), is_return=True, is_error=True, error_message=str(error),
-        trace_id=trace_id
+        source=target,
+        target=source,
+        action=action,
+        message=str(error),
+        is_return=True,
+        is_error=True,
+        error_message=str(error),
+        trace_id=trace_id,
     )
     logger.error(f"{target}-x{source}: Error", extra={"flow_event": err_event})
 
+
 @overload
-def trace_interaction(func: F) -> F:
-    ...
+def trace_interaction(func: F) -> F: ...
+
 
 @overload
 def trace_interaction(
-    *, 
-    source: Optional[str] = None, 
-    target: Optional[str] = None, 
-    action: Optional[str] = None
-) -> Callable[[F], F]:
-    ...
+    *,
+    source: Optional[str] = None,
+    target: Optional[str] = None,
+    name: Optional[str] = None,
+    action: Optional[str] = None,
+    capture_args: bool = True,
+    max_arg_length: int = 50,
+    max_arg_depth: int = 1,
+) -> Callable[[F], F]: ...
+
 
 def trace_interaction(
-    func: Optional[F] = None, 
-    *, 
-    source: Optional[str] = None, 
-    target: Optional[str] = None, 
-    action: Optional[str] = None
+    func: Optional[F] = None,
+    *,
+    source: Optional[str] = None,
+    target: Optional[str] = None,
+    name: Optional[str] = None,
+    action: Optional[str] = None,
+    capture_args: bool = True,
+    max_arg_length: int = 50,
+    max_arg_depth: int = 1,
 ) -> Union[F, Callable[[F], F]]:
     """
     Main Decorator for tracing function execution in Mermaid diagrams.
-    
+
     It supports two modes of operation:
     1. **Simple**: `@trace` (No arguments)
     2. **Configured**: `@trace(action="Login", target="AuthService")`
-    
+
     Args:
         func: The function being decorated (automatically passed in simple mode).
         source: Explicit name of the caller participant (rarely used, usually inferred from Context).
         target: Explicit name of the callee participant (overrides automatic resolution).
+        name: Alias for 'target' (for clearer API usage).
         action: Label for the arrow (defaults to function name).
+        capture_args: Whether to include arguments and return values in the log. Default True.
+        max_arg_length: Maximum string length for argument/result representation. Default 50.
+        max_arg_depth: Maximum recursion depth for argument/result representation. Default 1.
     """
-    
+
+    # Handle alias
+    final_target = target or name
+
     # Mode 1: @trace used without parentheses
     if func is not None and callable(func):
-        return _create_decorator(func, source, target, action)
-        
+        return _create_decorator(
+            func,
+            source,
+            final_target,
+            action,
+            capture_args,
+            max_arg_length,
+            max_arg_depth,
+        )
+
     # Mode 2: @trace(...) used with arguments -> returns a factory
     def factory(f: F) -> F:
-        return _create_decorator(f, source, target, action)
+        return _create_decorator(
+            f, source, final_target, action, capture_args, max_arg_length, max_arg_depth
+        )
+
     return factory
 
+
 def _create_decorator(
-    func: F, 
-    source: Optional[str], 
-    target: Optional[str], 
-    action: Optional[str]
+    func: F,
+    source: Optional[str],
+    target: Optional[str],
+    action: Optional[str],
+    capture_args: bool,
+    max_arg_length: int,
+    max_arg_depth: int,
 ) -> F:
     """
     Constructs the actual wrapper function.
     Handles both synchronous and asynchronous functions.
+
+    This function separates the wrapper creation logic from the argument parsing logic
+    in `trace_interaction`, making the code cleaner and easier to test.
     """
-    
+
     # Pre-calculate static metadata to save time at runtime
     if action is None:
+        # Default action name is the function name, converted to Title Case
         action = func.__name__.replace("_", " ").title()
 
     @functools.wraps(func)
@@ -197,49 +283,90 @@ def _create_decorator(
         """Sync function wrapper."""
         # 1. Resolve Context
         # 'source' is who called us (from Context). 'target' is who we are (resolved from self/cls).
+        # If 'source' is not explicitly provided, we look up the 'participant' set by the caller.
         current_source = source or LogContext.current_participant()
         trace_id = LogContext.current_trace_id()
         current_target = _resolve_target(func, args, target)
-        
+
         logger = get_flow_logger()
-        params_str = _format_args(args, kwargs)
-        
+        # Format arguments for the diagram arrow label
+        params_str = _format_args(
+            args, kwargs, capture_args, max_arg_length, max_arg_depth
+        )
+
         # 2. Log Request (Start of block)
-        _log_interaction(logger, current_source, current_target, action, params_str, trace_id)
-        
+        # Logs the initial "Call" arrow (Source -> Target)
+        _log_interaction(
+            logger, current_source, current_target, action, params_str, trace_id
+        )
+
         # 3. Execute with New Context
         # We push 'current_target' as the NEW 'participant' (source) for any internal calls.
+        # This builds the chain: A -> B, then inside B, B becomes the source for C (B -> C).
         with LogContext.scope({"participant": current_target, "trace_id": trace_id}):
             try:
                 result = func(*args, **kwargs)
                 # 4. Log Success Return
-                _log_return(logger, current_source, current_target, action, result, trace_id)
+                # Logs the "Return" arrow (Target --> Source)
+                _log_return(
+                    logger,
+                    current_source,
+                    current_target,
+                    action,
+                    result,
+                    trace_id,
+                    capture_args,
+                    max_arg_length,
+                    max_arg_depth,
+                )
                 return result
             except Exception as e:
                 # 5. Log Error Return
+                # Logs the "Error" arrow (Target --x Source)
                 _log_error(logger, current_source, current_target, action, e, trace_id)
                 raise
 
     @functools.wraps(func)
     async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
         """Async function wrapper (coroutine)."""
+        # 1. Resolve Context (Same as sync)
         current_source = source or LogContext.current_participant()
         trace_id = LogContext.current_trace_id()
         current_target = _resolve_target(func, args, target)
-        
+
         logger = get_flow_logger()
-        params_str = _format_args(args, kwargs)
-        
-        # 2. Log Request (Start of block)
-        _log_interaction(logger, current_source, current_target, action, params_str, trace_id)
-        
-        # Use async context manager (ascope) to ensure context propagates correctly across awaits
-        async with LogContext.ascope({"participant": current_target, "trace_id": trace_id}):
+        params_str = _format_args(
+            args, kwargs, capture_args, max_arg_length, max_arg_depth
+        )
+
+        # 2. Log Request
+        _log_interaction(
+            logger, current_source, current_target, action, params_str, trace_id
+        )
+
+        # 3. Execute with New Context using 'ascope'
+        # Use async context manager (ascope) to ensure context propagates correctly across awaits.
+        # This is critical for asyncio: context must be preserved even if the task yields control.
+        async with LogContext.ascope(
+            {"participant": current_target, "trace_id": trace_id}
+        ):
             try:
                 result = await func(*args, **kwargs)
-                _log_return(logger, current_source, current_target, action, result, trace_id)
+                # 4. Log Success Return
+                _log_return(
+                    logger,
+                    current_source,
+                    current_target,
+                    action,
+                    result,
+                    trace_id,
+                    capture_args,
+                    max_arg_length,
+                    max_arg_depth,
+                )
                 return result
             except Exception as e:
+                # 5. Log Error Return
                 _log_error(logger, current_source, current_target, action, e, trace_id)
                 raise
 
@@ -248,5 +375,6 @@ def _create_decorator(
         return cast(F, async_wrapper)
     return cast(F, wrapper)
 
+
 # Alias for easy import
 trace = trace_interaction

mermaid_trace/core/events.py

@@ -2,66 +2,68 @@ from dataclasses import dataclass, field
 import time
 from typing import Optional
 
+
 @dataclass
 class FlowEvent:
     """
     Represents a single interaction or step in the execution flow.
-    
+
     This data structure acts as the intermediate representation (IR) between
     runtime code execution and the final Mermaid diagram output. Each instance
     corresponds directly to one arrow or note in the sequence diagram.
-    
+
     The fields map to Mermaid syntax components as follows:
     `source` -> `target`: `message`
-    
+
     Attributes:
-        source (str): 
+        source (str):
             The name of the participant initiating the action (the "Caller").
             In Mermaid: The participant on the LEFT side of the arrow.
-            
-        target (str): 
+
+        target (str):
             The name of the participant receiving the action (the "Callee").
             In Mermaid: The participant on the RIGHT side of the arrow.
-            
-        action (str): 
+
+        action (str):
             A short, human-readable name for the operation (e.g., function name).
             Used for grouping or filtering logs, but often redundant with message.
-            
-        message (str): 
+
+        message (str):
             The actual text label displayed on the diagram arrow.
             Example: "getUser(id=1)" or "Return: User(name='Alice')".
-            
-        timestamp (float): 
+
+        timestamp (float):
             Unix timestamp (seconds) of when the event occurred.
-            Used for ordering events if logs are processed asynchronously, 
+            Used for ordering events if logs are processed asynchronously,
             though Mermaid sequence diagrams primarily rely on line order.
-            
-        trace_id (str): 
+
+        trace_id (str):
             Unique identifier for the trace session.
             Allows filtering multiple concurrent traces from a single log file
             to generate separate diagrams for separate requests.
-            
-        is_return (bool): 
+
+        is_return (bool):
             Flag indicating if this is a response arrow.
             If True, the arrow is drawn as a dotted line (`-->`) in Mermaid.
             If False, it is a solid line (`->`) representing a call.
-            
-        is_error (bool): 
+
+        is_error (bool):
             Flag indicating if an exception occurred.
             If True, the arrow might be styled differently (e.g., `-x`) to show failure.
-            
-        error_message (Optional[str]): 
+
+        error_message (Optional[str]):
             Detailed error text if `is_error` is True.
             Can be added as a note or included in the arrow label.
-            
-        params (Optional[str]): 
+
+        params (Optional[str]):
             Stringified representation of function arguments.
             Captured only for request events (call start).
-            
-        result (Optional[str]): 
+
+        result (Optional[str]):
             Stringified representation of the return value.
             Captured only for return events (call end).
     """
+
     source: str
     target: str
     action: str

mermaid_trace/core/formatter.py

@@ -3,6 +3,7 @@ import re
 from typing import Optional
 from .events import FlowEvent
 
+
 class MermaidFormatter(logging.Formatter):
     """
     Custom formatter to convert FlowEvents into Mermaid sequence diagram syntax.
@@ -10,8 +11,8 @@ class MermaidFormatter(logging.Formatter):
 
     def format(self, record: logging.LogRecord) -> str:
         # 1. Retrieve the FlowEvent
-        event: Optional[FlowEvent] = getattr(record, 'flow_event', None)
-        
+        event: Optional[FlowEvent] = getattr(record, "flow_event", None)
+
         if not event:
             # Fallback for standard logs if they accidentally reach this handler
             return super().format(record)
@@ -23,39 +24,47 @@ class MermaidFormatter(logging.Formatter):
         """
         Converts a FlowEvent into a Mermaid syntax string.
         """
-        # Sanitize participant names
+        # Sanitize participant names to avoid syntax errors in Mermaid
         src = self._sanitize(event.source)
         tgt = self._sanitize(event.target)
-        
+
         # Determine arrow type
         # ->> : Solid line with arrowhead (synchronous call)
         # -->> : Dotted line with arrowhead (return)
         # --x : Dotted line with cross (error)
         arrow = "-->>" if event.is_return else "->>"
-        
+
+        msg = ""
         if event.is_error:
             arrow = "--x"
             msg = f"Error: {event.error_message}"
         elif event.is_return:
+            # For returns, we usually show the return value or just "Return"
             msg = f"Return: {event.result}" if event.result else "Return"
         else:
+            # For calls, we show Action(Params) or just Action
             msg = f"{event.message}({event.params})" if event.params else event.message
-            
+
         # Optional: Add note or group if trace_id changes (not implemented in single line format)
         # For now, we just output the interaction.
-        
-        # Escape message for Mermaid safety
+
+        # Escape message for Mermaid safety (e.g. replacing newlines)
         msg = self._escape_message(msg)
-        
+
+        # Format: Source->>Target: Message
         return f"{src}{arrow}{tgt}: {msg}"
 
     def _sanitize(self, name: str) -> str:
         """
         Sanitizes participant names to be valid Mermaid identifiers.
         Allows alphanumeric and underscores. Replaces others.
+
+        Mermaid doesn't like spaces or special characters in participant aliases
+        unless they are quoted (which we are not doing here for simplicity),
+        so we replace them with underscores.
         """
         # Replace any non-alphanumeric character (except underscore) with underscore
-        clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
+        clean_name = re.sub(r"[^a-zA-Z0-9_]", "_", name)
         # Ensure it doesn't start with a digit (Mermaid doesn't like that sometimes, though often okay)
         if clean_name and clean_name[0].isdigit():
             clean_name = "_" + clean_name
@@ -67,6 +76,6 @@ class MermaidFormatter(logging.Formatter):
         Mermaid messages can contain most chars, but : and newlines can be tricky.
         """
         # Replace newlines with <br/> for Mermaid display
-        msg = msg.replace('\n', '<br/>')
+        msg = msg.replace("\n", "<br/>")
         # We might want to escape other chars if needed, but usually text after : is forgiving.
         return msg

mermaid_trace/handlers/async_handler.py

@@ -0,0 +1,52 @@
+import logging
+import logging.handlers
+import queue
+import atexit
+from typing import List, Optional
+
+
+class AsyncMermaidHandler(logging.handlers.QueueHandler):
+    """
+    A non-blocking logging handler that uses a background thread to write logs.
+
+    This handler pushes log records to a queue, which are then picked up by a
+    QueueListener running in a separate thread and dispatched to the actual
+    handlers (e.g., MermaidFileHandler).
+    """
+
+    def __init__(self, handlers: List[logging.Handler], queue_size: int = -1):
+        """
+        Initialize the async handler.
+        
+        Args:
+            handlers: A list of handlers that should receive the logs from the queue.
+                      (e.g., [MermaidFileHandler(...)])
+            queue_size: The maximum size of the queue. -1 means infinite.
+        """
+        self._log_queue: queue.Queue[logging.LogRecord] = queue.Queue(queue_size)
+        super().__init__(self._log_queue)
+        
+        # Initialize QueueListener
+        # It starts an internal thread to monitor the queue
+        # respect_handler_level=True ensures the target handlers' log levels are respected
+        self._listener: Optional[logging.handlers.QueueListener] = logging.handlers.QueueListener(
+            self._log_queue, 
+            *handlers, 
+            respect_handler_level=True
+        )
+        self._listener.start()
+        
+        # Ensure the listener is stopped and queue is flushed upon exit
+        # This prevents lost logs at program termination
+        atexit.register(self.stop)
+
+    def stop(self) -> None:
+        """
+        Stops the listener and flushes the queue.
+
+        This is registered with `atexit` to ensure that all pending logs
+        are written to disk before the application terminates.
+        """
+        if self._listener:
+            self._listener.stop()
+            self._listener = None