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

mermaid_trace/core/decorators.py

@@ -1,3 +1,15 @@
+"""
+Function Tracing Decorator Module
+
+This module provides the core tracing functionality for MermaidTrace. It includes:
+- A decorator (`trace`/`trace_interaction`) to instrument functions for tracing
+- Helper functions for formatting and logging trace events
+- Context management for tracking call chains
+
+The decorator can be applied to both synchronous and asynchronous functions,
+and automatically handles context propagation for nested calls.
+"""
+
 import functools
 import logging
 import inspect
@@ -7,6 +19,7 @@ from typing import Optional, Any, Callable, Tuple, Dict, Union, TypeVar, cast, o
 from .events import FlowEvent
 from .context import LogContext
 
+# Logger name for flow events - used to isolate tracing logs from other application logs
 FLOW_LOGGER_NAME = "mermaid_trace.flow"
 
 # Define generic type variable for the decorated function
@@ -14,7 +27,11 @@ F = TypeVar("F", bound=Callable[..., Any])
 
 
 def get_flow_logger() -> logging.Logger:
-    """Returns the dedicated logger for flow events."""
+    """Returns the dedicated logger for flow events.
+
+    Returns:
+        logging.Logger: Logger instance configured for tracing events
+    """
     return logging.getLogger(FLOW_LOGGER_NAME)
 
 
@@ -24,6 +41,14 @@ def _safe_repr(obj: Any, max_len: int = 50, max_depth: int = 1) -> str:
 
     Prevents massive log files by truncating long strings/objects
     and handling exceptions during __repr__ calls (e.g. strict objects).
+
+    Args:
+        obj: The object to represent as a string
+        max_len: Maximum length of the resulting string
+        max_depth: Maximum recursion depth for nested objects
+
+    Returns:
+        str: Safe, truncated representation of the object
     """
     try:
         # Create a custom repr object to control depth and length
@@ -37,6 +62,7 @@ def _safe_repr(obj: Any, max_len: int = 50, max_depth: int = 1) -> str:
             return r[:max_len] + "..."
         return r
     except Exception:
+        # Fallback if repr() fails for any reason
         return "<unrepresentable>"
 
 
@@ -50,14 +76,27 @@ def _format_args(
     """
     Formats function arguments into a single string "arg1, arg2, k=v".
     Used for the arrow label in the diagram.
+
+    Args:
+        args: Positional arguments to format
+        kwargs: Keyword arguments to format
+        capture_args: Whether to capture and format arguments at all
+        max_arg_length: Maximum length of each argument representation
+        max_arg_depth: Maximum recursion depth for nested arguments
+
+    Returns:
+        str: Formatted string of arguments, or empty string if capture_args is False
     """
     if not capture_args:
         return ""
 
-    parts = []
+    parts: list[str] = []
+
+    # Format positional arguments
     for arg in args:
         parts.append(_safe_repr(arg, max_len=max_arg_length, max_depth=max_arg_depth))
 
+    # Format keyword arguments
     for k, v in kwargs.items():
         val_str = _safe_repr(v, max_len=max_arg_length, max_depth=max_arg_depth)
         parts.append(f"{k}={val_str}")
@@ -78,16 +117,22 @@ def _resolve_target(
     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.
     5. **Fallback**: "Unknown".
+
+    Args:
+        func: The function being called
+        args: Positional arguments passed to the function
+        target_override: Explicit target name provided by the user, if any
+
+    Returns:
+        str: Resolved target name for the diagram
     """
     if target_override:
         return target_override
 
-    # Heuristic: If it's a method call, args[0] is usually 'self'.
+    # Heuristic: If it's a method call, args[0] is usually 'self' or 'cls'
     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.
+        # Check if it looks like a class instance (not a primitive or container)
         if hasattr(first_arg, "__class__") and not isinstance(
             first_arg, (str, int, float, bool, list, dict, type)
         ):
@@ -113,7 +158,15 @@ def _log_interaction(
 ) -> None:
     """
     Logs the 'Call' event (Start of function).
-    Arrow: source -> target
+    Generates a FlowEvent and logs it with the appropriate context.
+
+    Args:
+        logger: Logger instance to use for logging
+        source: Name of the source participant (caller)
+        target: Name of the target participant (callee)
+        action: Name of the action being performed
+        params: Formatted string of function arguments
+        trace_id: Unique trace identifier for correlation
     """
     req_event = FlowEvent(
         source=source,
@@ -144,6 +197,17 @@ def _log_return(
 
     Note: 'source' here is the original caller, 'target' is the callee.
     So the return arrow goes from target back to source.
+
+    Args:
+        logger: Logger instance to use for logging
+        source: Name of the original caller
+        target: Name of the callee that's returning
+        action: Name of the action being returned from
+        result: Return value of the function
+        trace_id: Unique trace identifier for correlation
+        capture_args: Whether to include the return value in the log
+        max_arg_length: Maximum length of the return value representation
+        max_arg_depth: Maximum recursion depth for nested return values
     """
     result_str = ""
     if capture_args:
@@ -172,6 +236,14 @@ def _log_error(
     """
     Logs an 'Error' event if the function raises an exception.
     Arrow: target -x source (Error return)
+
+    Args:
+        logger: Logger instance to use for logging
+        source: Name of the original caller
+        target: Name of the callee that encountered an error
+        action: Name of the action that failed
+        error: Exception that was raised
+        trace_id: Unique trace identifier for correlation
     """
     err_event = FlowEvent(
         source=target,
@@ -217,25 +289,32 @@ def trace_interaction(
     """
     Main Decorator for tracing function execution in Mermaid diagrams.
 
+    This decorator instruments functions to log their execution flow as Mermaid
+    sequence diagram events. It supports both synchronous and asynchronous functions,
+    and automatically handles context propagation for nested calls.
+
     It supports two modes of operation:
-    1. **Simple**: `@trace` (No arguments)
-    2. **Configured**: `@trace(action="Login", target="AuthService")`
+    1. **Simple**: `@trace` (No arguments) - uses default settings
+    2. **Configured**: `@trace(action="Login", target="AuthService")` - customizes behavior
 
     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).
+        action: Label for the arrow (defaults to function name in Title Case).
         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.
+
+    Returns:
+        Callable: Either the decorated function (simple mode) or a decorator factory (configured mode)
     """
 
-    # Handle alias
+    # Handle alias - 'name' is an alternative name for 'target'
     final_target = target or name
 
-    # Mode 1: @trace used without parentheses
+    # Mode 1: @trace used without parentheses - directly decorate the function
     if func is not None and callable(func):
         return _create_decorator(
             func,
@@ -247,7 +326,7 @@ def trace_interaction(
             max_arg_depth,
         )
 
-    # Mode 2: @trace(...) used with arguments -> returns a factory
+    # Mode 2: @trace(...) used with arguments -> returns a factory that will decorate the function
     def factory(f: F) -> F:
         return _create_decorator(
             f, source, final_target, action, capture_args, max_arg_length, max_arg_depth
@@ -266,11 +345,23 @@ def _create_decorator(
     max_arg_depth: int,
 ) -> F:
     """
-    Constructs the actual wrapper function.
-    Handles both synchronous and asynchronous functions.
+    Constructs the actual wrapper function for the decorated function.
+    Handles both synchronous and asynchronous functions by creating the appropriate wrapper.
 
     This function separates the wrapper creation logic from the argument parsing logic
     in `trace_interaction`, making the code cleaner and easier to test.
+
+    Args:
+        func: The function to decorate
+        source: Explicit source name, if any
+        target: Explicit target name, if any
+        action: Explicit action name, if any
+        capture_args: Whether to capture arguments and return values
+        max_arg_length: Maximum length for argument/return value representations
+        max_arg_depth: Maximum recursion depth for nested objects
+
+    Returns:
+        Callable: Decorated function with tracing logic
     """
 
     # Pre-calculate static metadata to save time at runtime
@@ -280,10 +371,9 @@ def _create_decorator(
 
     @functools.wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
-        """Sync function wrapper."""
+        """Synchronous function wrapper that adds tracing logic."""
         # 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)
@@ -294,7 +384,7 @@ def _create_decorator(
             args, kwargs, capture_args, max_arg_length, max_arg_depth
         )
 
-        # 2. Log Request (Start of block)
+        # 2. Log Request (Start of function)
         # Logs the initial "Call" arrow (Source -> Target)
         _log_interaction(
             logger, current_source, current_target, action, params_str, trace_id
@@ -328,7 +418,7 @@ def _create_decorator(
 
     @functools.wraps(func)
     async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-        """Async function wrapper (coroutine)."""
+        """Asynchronous function wrapper that adds tracing logic."""
         # 1. Resolve Context (Same as sync)
         current_source = source or LogContext.current_participant()
         trace_id = LogContext.current_trace_id()
@@ -372,9 +462,9 @@ def _create_decorator(
 
     # Detect if the wrapped function is a coroutine to choose the right wrapper
     if inspect.iscoroutinefunction(func):
-        return cast(F, async_wrapper)
-    return cast(F, wrapper)
+        return cast(F, async_wrapper)  # Return async wrapper for coroutines
+    return cast(F, wrapper)  # Return sync wrapper for regular functions
 
 
-# Alias for easy import
+# Alias for easy import - 'trace' is the primary name users should use
 trace = trace_interaction

mermaid_trace/core/events.py

@@ -1,28 +1,107 @@
+"""
+Event Definition Module
+
+This module defines the event structure for the tracing system. It provides an
+abstract Event base class and a concrete FlowEvent implementation that represents
+individual interactions in the execution flow.
+"""
+
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 import time
 from typing import Optional
 
 
+class Event(ABC):
+    """
+    Abstract base class for all event types.
+
+    This provides a common interface for different types of events, allowing
+    for extensibility and supporting multiple output formats. Concrete event
+    classes must implement all abstract methods.
+    """
+
+    @abstractmethod
+    def get_source(self) -> str:
+        """
+        Get the source of the event.
+
+        Returns:
+            str: Name of the participant that generated the event
+        """
+        pass
+
+    @abstractmethod
+    def get_target(self) -> str:
+        """
+        Get the target of the event.
+
+        Returns:
+            str: Name of the participant that received the event
+        """
+        pass
+
+    @abstractmethod
+    def get_action(self) -> str:
+        """
+        Get the action name of the event.
+
+        Returns:
+            str: Short name describing the action performed
+        """
+        pass
+
+    @abstractmethod
+    def get_message(self) -> str:
+        """
+        Get the message text of the event.
+
+        Returns:
+            str: Detailed message describing the event
+        """
+        pass
+
+    @abstractmethod
+    def get_timestamp(self) -> float:
+        """
+        Get the timestamp of the event.
+
+        Returns:
+            float: Unix timestamp (seconds) when the event occurred
+        """
+        pass
+
+    @abstractmethod
+    def get_trace_id(self) -> str:
+        """
+        Get the trace ID of the event.
+
+        Returns:
+            str: Unique identifier for the trace session
+        """
+        pass
+
+
 @dataclass
-class FlowEvent:
+class FlowEvent(Event):
     """
     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
+    runtime code execution and the final diagram output. Each instance
     corresponds directly to one arrow or note in the sequence diagram.
 
-    The fields map to Mermaid syntax components as follows:
+    The fields map to diagram syntax components as follows:
     `source` -> `target`: `message`
 
     Attributes:
         source (str):
             The name of the participant initiating the action (the "Caller").
-            In Mermaid: The participant on the LEFT side of the arrow.
+            In sequence diagrams: The participant on the LEFT side of the arrow.
 
         target (str):
             The name of the participant receiving the action (the "Callee").
-            In Mermaid: The participant on the RIGHT side of the arrow.
+            In sequence diagrams: The participant on the RIGHT side of the arrow.
 
         action (str):
             A short, human-readable name for the operation (e.g., function name).
@@ -34,8 +113,8 @@ class FlowEvent:
 
         timestamp (float):
             Unix timestamp (seconds) of when the event occurred.
-            Used for ordering events if logs are processed asynchronously,
-            though Mermaid sequence diagrams primarily rely on line order.
+            Used for ordering events if logs are processed asynchronously.
+            Defaults to current time when event is created.
 
         trace_id (str):
             Unique identifier for the trace session.
@@ -44,34 +123,68 @@ class FlowEvent:
 
         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.
+            If True, the arrow is drawn as a dotted line in sequence diagrams.
+            If False, it is a solid line representing a call.
+            Defaults to False.
 
         is_error (bool):
             Flag indicating if an exception occurred.
-            If True, the arrow might be styled differently (e.g., `-x`) to show failure.
+            If True, the arrow might be styled differently to show failure.
+            Defaults to False.
 
         error_message (Optional[str]):
             Detailed error text if `is_error` is True.
             Can be added as a note or included in the arrow label.
+            Defaults to None.
 
         params (Optional[str]):
             Stringified representation of function arguments.
             Captured only for request events (call start).
+            Defaults to None.
 
         result (Optional[str]):
             Stringified representation of the return value.
             Captured only for return events (call end).
+            Defaults to None.
     """
 
-    source: str
-    target: str
-    action: str
-    message: str
-    trace_id: str
-    timestamp: float = field(default_factory=time.time)
-    is_return: bool = False
-    is_error: bool = False
-    error_message: Optional[str] = None
-    params: Optional[str] = None
-    result: Optional[str] = None
+    # Required fields for every event
+    source: str  # Participant who initiated the action
+    target: str  # Participant who received the action
+    action: str  # Short name for the operation
+    message: str  # Detailed message for the diagram arrow
+    trace_id: str  # Unique identifier for the trace session
+
+    # Optional fields with defaults
+    timestamp: float = field(
+        default_factory=time.time
+    )  # Unix timestamp of event creation
+    is_return: bool = False  # Whether this is a response arrow
+    is_error: bool = False  # Whether an error occurred
+    error_message: Optional[str] = None  # Detailed error message if is_error is True
+    params: Optional[str] = None  # Stringified function arguments
+    result: Optional[str] = None  # Stringified return value
+
+    def get_source(self) -> str:
+        """Get the source of the event."""
+        return self.source
+
+    def get_target(self) -> str:
+        """Get the target of the event."""
+        return self.target
+
+    def get_action(self) -> str:
+        """Get the action name of the event."""
+        return self.action
+
+    def get_message(self) -> str:
+        """Get the message text of the event."""
+        return self.message
+
+    def get_timestamp(self) -> float:
+        """Get the timestamp of the event."""
+        return self.timestamp
+
+    def get_trace_id(self) -> str:
+        """Get the trace ID of the event."""
+        return self.trace_id

mermaid_trace/core/formatter.py

@@ -1,81 +1,151 @@
+"""
+Event Formatting Module
+
+This module provides formatters to convert Event objects into various output formats.
+Currently, it supports Mermaid sequence diagram syntax formatting, but can be extended
+with additional formatters for other diagram types or logging formats.
+"""
+
 import logging
 import re
+from abc import ABC, abstractmethod
 from typing import Optional
-from .events import FlowEvent
+from .events import Event, FlowEvent
 
 
-class MermaidFormatter(logging.Formatter):
+class BaseFormatter(ABC, logging.Formatter):
     """
-    Custom formatter to convert FlowEvents into Mermaid sequence diagram syntax.
+    Abstract base class for all event formatters.
+
+    This provides a common interface for different formatters, allowing
+    for extensibility and supporting multiple output formats.
+
+    Subclasses must implement the format_event method to convert Event objects
+    into the desired output string format.
     """
 
+    @abstractmethod
+    def format_event(self, event: Event) -> str:
+        """
+        Format an Event into the desired output string.
+
+        Args:
+            event: The Event object to format
+
+        Returns:
+            str: Formatted string representation of the event
+        """
+        pass
+
     def format(self, record: logging.LogRecord) -> str:
-        # 1. Retrieve the FlowEvent
-        event: Optional[FlowEvent] = getattr(record, "flow_event", None)
+        """
+        Format a logging record containing an event.
+
+        This method overrides the standard logging.Formatter.format method to
+        extract and format Event objects from log records.
+
+        Args:
+            record: The logging record to format. Must contain a 'flow_event' attribute
+                   if it represents a tracing event.
+
+        Returns:
+            str: Formatted string representation of the record
+        """
+        # Retrieve the Event object from the log record
+        event: Optional[Event] = getattr(record, "flow_event", None)
 
         if not event:
             # Fallback for standard logs if they accidentally reach this handler
             return super().format(record)
 
-        # 2. Convert event to Mermaid line
-        return self._to_mermaid_line(event)
+        # Convert event to the desired format using the subclass's format_event method
+        return self.format_event(event)
 
-    def _to_mermaid_line(self, event: FlowEvent) -> str:
+
+class MermaidFormatter(BaseFormatter):
+    """
+    Custom formatter to convert Events into Mermaid sequence diagram syntax.
+
+    This formatter transforms FlowEvent objects into lines of Mermaid syntax that
+    can be directly written to a .mmd file. Each event becomes a single line in the
+    sequence diagram.
+    """
+
+    def format_event(self, event: Event) -> str:
         """
-        Converts a FlowEvent into a Mermaid syntax string.
+        Converts an Event into a Mermaid syntax string.
+
+        Args:
+            event: The Event object to format
+
+        Returns:
+            str: Mermaid syntax string representation of the event
         """
+        if not isinstance(event, FlowEvent):
+            # Fallback format for non-FlowEvent types
+            return f"{event.get_source()}->>{event.get_target()}: {event.get_message()}"
+
         # Sanitize participant names to avoid syntax errors in Mermaid
         src = self._sanitize(event.source)
         tgt = self._sanitize(event.target)
 
-        # Determine arrow type
+        # Determine arrow type based on event properties
         # ->> : Solid line with arrowhead (synchronous call)
         # -->> : Dotted line with arrowhead (return)
         # --x : Dotted line with cross (error)
         arrow = "-->>" if event.is_return else "->>"
 
+        # Construct message text based on event type
         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"
+            # For return events, show 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
+            # For call events, 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 (e.g. replacing newlines)
+        # Escape message for Mermaid safety (e.g., replacing newlines)
         msg = self._escape_message(msg)
 
+        # Return the complete Mermaid syntax line
         # 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.
+
+        Args:
+            name: Original participant name
+
+        Returns:
+            str: Sanitized participant name
         """
         # Replace any non-alphanumeric character (except underscore) with underscore
         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)
+        # Ensure it doesn't start with a digit (Mermaid doesn't like that sometimes)
         if clean_name and clean_name[0].isdigit():
             clean_name = "_" + clean_name
         return clean_name
 
     def _escape_message(self, msg: str) -> str:
         """
-        Escapes special characters in the message text.
-        Mermaid messages can contain most chars, but : and newlines can be tricky.
+        Escapes special characters in the message text for safe Mermaid rendering.
+
+        Args:
+            msg: Original message text
+
+        Returns:
+            str: Escaped message text
         """
-        # Replace newlines with <br/> for Mermaid display
+        # Replace newlines with <br/> for proper display in Mermaid diagrams
         msg = msg.replace("\n", "<br/>")
-        # We might want to escape other chars if needed, but usually text after : is forgiving.
+        # Additional escaping could be added here if needed for other characters
         return msg