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

mermaid_trace/__init__.py

@@ -7,14 +7,28 @@ developers understand the flow of their applications, debug complex interactions
 and document system behavior.
 
 Key Components:
-- `trace`: A decorator to instrument functions for tracing.
+- `trace`: A decorator to instrument functions for tracing. It captures arguments,
+  return values, and errors, and logs them as interactions.
 - `LogContext`: Manages execution context (like thread-local storage) to track
   caller/callee relationships across async tasks and threads.
 - `configure_flow`: Sets up the logging handler to write diagrams to a file.
+  It handles handler configuration, file modes, and async logging setup.
+
+Usage Example:
+    from mermaid_trace import trace, configure_flow
+
+    configure_flow("my_flow.mmd")
+
+    @trace
+    def hello():
+        print("Hello")
+
+    hello()
 """
 
 from .core.decorators import trace_interaction, trace
 from .handlers.mermaid_handler import MermaidFileHandler
+from .handlers.async_handler import AsyncMermaidHandler
 from .core.events import FlowEvent
 from .core.context import LogContext
 from .core.formatter import MermaidFormatter
@@ -22,42 +36,80 @@ from .core.formatter import MermaidFormatter
 # Integrations (like FastAPI) must be imported explicitly by the user if needed.
 
 from importlib.metadata import PackageNotFoundError, version
+from typing import List, Optional
 
 import logging
 
-def configure_flow(output_file: str = "flow.mmd") -> logging.Logger:
+
+def configure_flow(
+    output_file: str = "flow.mmd",
+    handlers: Optional[List[logging.Handler]] = None,
+    append: bool = False,
+    async_mode: bool = False,
+) -> logging.Logger:
     """
     Configures the flow logger to output to a Mermaid file.
-    
+
     This function sets up the logging infrastructure required to capture
     trace events and write them to the specified output file. It should
-    be called once at the start of your application.
-    
+    be called once at the start of your application to initialize the tracing system.
+
     Args:
         output_file (str): The absolute or relative path to the output .mmd file.
                            Defaults to "flow.mmd" in the current directory.
-    
+                           If the file does not exist, it will be created with the correct header.
+        handlers (List[logging.Handler], optional): A list of custom logging handlers.
+                                                    If provided, 'output_file' is ignored unless
+                                                    you explicitly include a MermaidFileHandler.
+                                                    Useful if you want to stream logs to other destinations.
+        append (bool): If True, adds the new handler(s) without removing existing ones.
+                       Defaults to False (clears existing handlers to prevent duplicate logging).
+        async_mode (bool): If True, uses a non-blocking background thread for logging (QueueHandler).
+                           Recommended for high-performance production environments to avoid
+                           blocking the main execution thread during file I/O.
+                           Defaults to False.
+
     Returns:
         logging.Logger: The configured logger instance used for flow tracing.
     """
     # Get the specific logger used by the tracing decorators
+    # This logger is isolated from the root logger to prevent pollution
     logger = logging.getLogger("mermaid_trace.flow")
     logger.setLevel(logging.INFO)
-    
+
     # Remove existing handlers to avoid duplicate logs if configured multiple times
-    if logger.hasHandlers():
+    # unless 'append' is requested. This ensures idempotency when calling configure_flow multiple times.
+    if not append and logger.hasHandlers():
         logger.handlers.clear()
-        
-    # Create and attach the custom handler that writes Mermaid syntax
-    handler = MermaidFileHandler(output_file)
-    
-    # Use the custom formatter to convert FlowEvents to Mermaid strings
-    handler.setFormatter(MermaidFormatter())
-    
-    logger.addHandler(handler)
-    
+
+    # Determine the target handlers
+    target_handlers = []
+
+    if handlers:
+        # Use user-provided handlers if available
+        target_handlers = handlers
+    else:
+        # Create default Mermaid handler
+        # This handler knows how to write the Mermaid header and format events
+        handler = MermaidFileHandler(output_file)
+        handler.setFormatter(MermaidFormatter())
+        target_handlers = [handler]
+
+    if async_mode:
+        # Wrap the target handlers in an AsyncMermaidHandler (QueueHandler)
+        # The QueueListener will pick up logs from the queue and dispatch to target_handlers
+        # This decouples the application execution from the logging I/O
+        async_handler = AsyncMermaidHandler(target_handlers)
+        logger.addHandler(async_handler)
+    else:
+        # Attach handlers directly to the logger for synchronous logging
+        # Simple and reliable for debugging or low-throughput applications
+        for h in target_handlers:
+            logger.addHandler(h)
+
     return logger
 
+
 try:
     # Attempt to retrieve the installed package version
     __version__ = version("mermaid-trace")
@@ -67,4 +119,13 @@ except PackageNotFoundError:
 
 
 # Export public API for easy access
-__all__ = ["trace_interaction", "trace", "configure_flow", "MermaidFileHandler", "LogContext", "FlowEvent", "MermaidFormatter"]
+__all__ = [
+    "trace_interaction",
+    "trace",
+    "configure_flow",
+    "MermaidFileHandler",
+    "AsyncMermaidHandler",
+    "LogContext",
+    "FlowEvent",
+    "MermaidFormatter",
+]

mermaid_trace/cli.py

@@ -8,13 +8,22 @@ from pathlib import Path
 from typing import Type, Any
 
 try:
-    from watchdog.observers import Observer  
-    from watchdog.events import FileSystemEventHandler  
+    # Watchdog is an optional dependency that allows efficient file monitoring.
+    # If installed, we use it to detect file changes instantly.
+    from watchdog.observers import Observer
+    from watchdog.events import FileSystemEventHandler
+
     HAS_WATCHDOG = True
 except ImportError:
+    # Fallback for when watchdog is not installed (e.g., minimal install).
     HAS_WATCHDOG = False
 
 # HTML Template for the preview page
+# This template provides a self-contained environment to render Mermaid diagrams.
+# It includes:
+# 1. Mermaid.js library from CDN.
+# 2. CSS for basic styling and layout.
+# 3. JavaScript logic for auto-refreshing when the source file changes.
 HTML_TEMPLATE = """
 <!DOCTYPE html>
 <html lang="en">
@@ -59,12 +68,15 @@ HTML_TEMPLATE = """
         mermaid.initialize({{ startOnLoad: true, theme: 'default' }});
         
         // Live Reload Logic
+        // We track the file's modification time (mtime) sent by the server.
         const currentMtime = "{mtime}";
         
         function checkUpdate() {{
+            // Poll the /_status endpoint to check if the file has changed on disk
             fetch('/_status')
                 .then(response => response.text())
                 .then(mtime => {{
+                    // If the server reports a different mtime, reload the page
                     if (mtime && mtime !== currentMtime) {{
                         console.log("File changed, reloading...");
                         location.reload();
@@ -74,20 +86,40 @@ HTML_TEMPLATE = """
         }}
         
         // Poll every 1 second
+        // This is a simple alternative to WebSockets for local dev tools
         setInterval(checkUpdate, 1000);
     </script>
 </body>
 </html>
 """
 
-def _create_handler(filename: str, path: Path) -> Type[http.server.SimpleHTTPRequestHandler]:
-    """Helper to create the request handler class with closure over filename/path."""
+
+def _create_handler(
+    filename: str, path: Path
+) -> Type[http.server.SimpleHTTPRequestHandler]:
+    """
+    Factory function to create a custom request handler class.
+
+    This uses a closure to inject `filename` and `path` into the handler's scope,
+    allowing the `do_GET` method to access them without global variables.
+
+    Args:
+        filename (str): Display name of the file.
+        path (Path): Path object to the file on disk.
+
+    Returns:
+        Type[SimpleHTTPRequestHandler]: A custom handler class.
+    """
+
     class Handler(http.server.SimpleHTTPRequestHandler):
         """
         Custom Request Handler to serve the generated HTML dynamically.
+        It intercepts GET requests to serve the constructed HTML instead of static files.
         """
+
         def log_message(self, format: str, *args: Any) -> None:
             # Suppress default logging to keep console clean
+            # We only want to see application logs, not every HTTP request
             pass
 
         def do_GET(self) -> None:
@@ -99,22 +131,26 @@ def _create_handler(filename: str, path: Path) -> Type[http.server.SimpleHTTPReq
                 self.send_response(200)
                 self.send_header("Content-type", "text/html")
                 self.end_headers()
-                
+
                 try:
                     # Read the current content of the mermaid file
                     content = path.read_text(encoding="utf-8")
                     mtime = str(path.stat().st_mtime)
                 except Exception as e:
                     # Fallback if reading fails (e.g., file locked)
+                    # Show the error directly in the diagram area
                     content = f"sequenceDiagram\nNote right of Error: Failed to read file: {e}"
                     mtime = "0"
-                
+
                 # Inject content into the HTML template
-                html = HTML_TEMPLATE.format(filename=filename, content=content, mtime=mtime)
+                html = HTML_TEMPLATE.format(
+                    filename=filename, content=content, mtime=mtime
+                )
                 self.wfile.write(html.encode("utf-8"))
-            
+
             elif self.path == "/_status":
-                # API endpoint for client to check file status
+                # API endpoint for client-side polling.
+                # Returns the current modification time of the file.
                 self.send_response(200)
                 self.send_header("Content-type", "text/plain")
                 self.end_headers()
@@ -123,15 +159,29 @@ def _create_handler(filename: str, path: Path) -> Type[http.server.SimpleHTTPReq
                 except OSError:
                     mtime = "0"
                 self.wfile.write(mtime.encode("utf-8"))
-                
+
             else:
-                # Serve static files if needed, or return 404
+                # Serve other static files if needed, or return 404
                 super().do_GET()
+
     return Handler
 
+
 def serve(filename: str, port: int = 8000) -> None:
     """
     Starts a local HTTP server to preview the Mermaid diagram.
+
+    This function blocks the main thread and runs a TCP server.
+    It automatically opens the default web browser to the preview URL.
+
+    Features:
+    - Serves the .mmd file wrapped in an HTML viewer.
+    - Uses Watchdog (if available) or client-side polling for live reloads.
+    - Gracefully handles shutdown on Ctrl+C.
+
+    Args:
+        filename (str): Path to the .mmd file to serve.
+        port (int): Port to bind the server to. Default is 8000.
     """
     path = Path(filename)
     if not path.exists():
@@ -139,11 +189,18 @@ def serve(filename: str, port: int = 8000) -> None:
         sys.exit(1)
 
     # Setup Watchdog if available
+    # Watchdog allows us to print console messages when the file changes.
+    # The actual browser reload is triggered by the client polling the /_status endpoint,
+    # but Watchdog gives immediate feedback in the terminal.
     observer = None
     if HAS_WATCHDOG:
+
         class FileChangeHandler(FileSystemEventHandler):
             def on_modified(self, event: Any) -> None:
-                if not event.is_directory and os.path.abspath(event.src_path) == str(path.resolve()):
+                # Filter for the specific file we are watching
+                if not event.is_directory and os.path.abspath(event.src_path) == str(
+                    path.resolve()
+                ):
                     print(f"[Watchdog] File changed: {filename}")
 
         print("Initializing file watcher...")
@@ -151,17 +208,21 @@ def serve(filename: str, port: int = 8000) -> None:
         observer.schedule(FileChangeHandler(), path=str(path.parent), recursive=False)
         observer.start()
     else:
-        print("Watchdog not installed. Falling back to polling mode (client-side only).")
+        print(
+            "Watchdog not installed. Falling back to polling mode (client-side only)."
+        )
 
     HandlerClass = _create_handler(filename, path)
 
     print(f"Serving {filename} at http://localhost:{port}")
     print("Press Ctrl+C to stop.")
-    
+
     # Open browser automatically to the server URL
     webbrowser.open(f"http://localhost:{port}")
-    
+
     # Start the TCP server
+    # ThreadingTCPServer is used to handle multiple requests concurrently if needed,
+    # ensuring the browser polling doesn't block the initial load.
     with socketserver.ThreadingTCPServer(("", port), HandlerClass) as httpd:
         try:
             httpd.serve_forever()
@@ -172,22 +233,29 @@ def serve(filename: str, port: int = 8000) -> None:
                 observer.join()
             httpd.server_close()
 
+
 def main() -> None:
     """
     Entry point for the CLI application.
+    Parses arguments and dispatches to the appropriate command handler.
     """
     parser = argparse.ArgumentParser(description="MermaidTrace CLI")
     subparsers = parser.add_subparsers(dest="command", required=True)
-    
+
     # 'serve' command definition
-    serve_parser = subparsers.add_parser("serve", help="Serve a Mermaid file in the browser")
+    serve_parser = subparsers.add_parser(
+        "serve", help="Serve a Mermaid file in the browser"
+    )
     serve_parser.add_argument("file", help="Path to the .mmd file")
-    serve_parser.add_argument("--port", type=int, default=8000, help="Port to bind to (default: 8000)")
-    
+    serve_parser.add_argument(
+        "--port", type=int, default=8000, help="Port to bind to (default: 8000)"
+    )
+
     args = parser.parse_args()
-    
+
     if args.command == "serve":
         serve(args.file, args.port)
 
+
 if __name__ == "__main__":
     main()

mermaid_trace/core/context.py

@@ -3,17 +3,18 @@ from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncIterator, Dict, Iterator
 import uuid
 
+
 class LogContext:
     """
     Manages global context information for logging (e.g., request_id, user_id, current_participant).
-    
-    This class utilizes `contextvars.ContextVar` to ensure thread-safety and 
+
+    This class utilizes `contextvars.ContextVar` to ensure thread-safety and
     correct context propagation in asynchronous (asyncio) environments.
-    Unlike `threading.local()`, `ContextVar` works natively with Python's async/await 
-    event loop, ensuring that context is preserved across `await` points but isolated 
+    Unlike `threading.local()`, `ContextVar` works natively with Python's async/await
+    event loop, ensuring that context is preserved across `await` points but isolated
     between different concurrent tasks.
     """
-    
+
     # ContextVar is the key mechanism here.
     # It stores a dictionary unique to the current execution context (Task/Thread).
     # "log_context" is the name of the variable, useful for debugging.
@@ -24,9 +25,9 @@ class LogContext:
     def _get_store(cls) -> Dict[str, Any]:
         """
         Retrieves the current context dictionary.
-        
-        If the context variable has not been set in the current context, 
-        it returns a fresh empty dictionary. This prevents LookupError 
+
+        If the context variable has not been set in the current context,
+        it returns a fresh empty dictionary. This prevents LookupError
         and ensures there's always a valid dictionary to work with.
         """
         try:
@@ -38,8 +39,8 @@ class LogContext:
     def set(cls, key: str, value: Any) -> None:
         """
         Sets a specific key-value pair in the current context.
-        
-        Important: ContextVars are immutable collections. To modify the context, 
+
+        Important: ContextVars are immutable collections. To modify the context,
         we must:
         1. Retrieve the current dictionary.
         2. Create a shallow copy (to avoid affecting parent contexts if we were reusing the object).
@@ -54,8 +55,8 @@ class LogContext:
     def update(cls, data: Dict[str, Any]) -> None:
         """
         Updates multiple keys in the current context at once.
-        
-        This follows the same Copy-Update-Set pattern as `set()` to maintain 
+
+        This follows the same Copy-Update-Set pattern as `set()` to maintain
         context isolation.
         """
         if not data:
@@ -83,18 +84,18 @@ class LogContext:
     def scope(cls, data: Dict[str, Any]) -> Iterator[None]:
         """
         Synchronous context manager for temporary context updates.
-        
+
         Usage:
             with LogContext.scope({"user_id": 123}):
                 # user_id is 123 here
                 some_function()
             # user_id reverts to previous value (or disappears) here
-            
+
         Mechanism:
             1. Copies current context and updates it with new data.
             2. Sets the ContextVar to this new state, receiving a `Token`.
             3. Yields control to the block.
-            4. Finally, uses the `Token` to reset the ContextVar to its exact state 
+            4. Finally, uses the `Token` to reset the ContextVar to its exact state
                before the block entered.
         """
         current_ctx = cls._get_store().copy()
@@ -111,11 +112,11 @@ class LogContext:
     async def ascope(cls, data: Dict[str, Any]) -> AsyncIterator[None]:
         """
         Async context manager for temporary context updates in coroutines.
-        
+
         Usage:
             async with LogContext.ascope({"request_id": "abc"}):
                 await some_async_function()
-                
+
         This is functionally identical to `scope` but designed for `async with` blocks.
         It ensures that even if the code inside `yield` suspends execution (await),
         the context remains valid for that task.
@@ -165,12 +166,12 @@ class LogContext:
     def current_trace_id(cls) -> str:
         """
         Retrieves the current trace ID for correlating events in a single flow.
-        
+
         Lazy Initialization Logic:
-        If no trace_id exists in the current context, it generates a new UUIDv4 
+        If no trace_id exists in the current context, it generates a new UUIDv4
         and sets it immediately. This ensures that:
         1. A trace ID is always available when asked for.
-        2. Once generated, the same ID persists for the duration of the context 
+        2. Once generated, the same ID persists for the duration of the context
            (unless manually changed), linking all subsequent logs together.
         """
         tid = cls.get("trace_id")