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

mermaid_trace/handlers/async_handler.py

@@ -1,105 +1,181 @@
 """
 Asynchronous Mermaid Handler Module
+===================================
 
-This module provides a non-blocking logging handler that uses a background thread
-for writing logs. It's designed to improve performance in high-throughput applications
-by decoupling the logging I/O from the main execution thread.
+This module implements a non-blocking logging handler that leverages a background thread
+and a queue mechanism to handle log records. This design pattern is crucial for
+high-performance applications where I/O latency (like writing to disk) in the main
+execution thread is unacceptable.
+
+Key Components:
+1.  **AsyncMermaidHandler**: The frontend handler that quickly pushes logs to a queue.
+2.  **Queue**: A thread-safe FIFO buffer (producer-consumer pattern).
+3.  **QueueListener**: A background worker that pulls logs from the queue and
+    writes them to the actual destination (e.g., a file).
+
+By decoupling log generation from log persistence, we ensure that the application's
+core logic remains responsive even during bursts of logging activity.
 """
 
 import logging
 import logging.handlers
 import queue
 import atexit
-from typing import List, Optional
+from typing import List, Optional, cast
 
 
 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).
-
-    This architecture provides several benefits:
-    - Main thread doesn't block waiting for disk I/O
-    - Logs are processed in the background
-    - Better performance in high-throughput applications
-    - Smooth handling of burst traffic
+    A high-performance, non-blocking logging handler using the Producer-Consumer pattern.
+
+    This handler acts as the "Producer". It intercepts log records and immediately
+    pushes them into a thread-safe queue. A separate "Consumer" thread (managed by
+    QueueListener) asynchronously picks up these records and dispatches them to the
+    actual underlying handlers (like MermaidFileHandler).
+
+    **Architecture & Benefits:**
+    - **Non-Blocking I/O**: The main application thread never waits for disk writes.
+      It only waits for the (very fast) queue insertion operation.
+    - **Burst Handling**: The queue acts as a buffer, absorbing sudden spikes in
+      log volume without slowing down the application.
+    - **Thread Safety**: Uses Python's thread-safe `queue.Queue` for synchronization.
+    - **Graceful Shutdown**: Integrates with `atexit` to ensure pending logs are
+      flushed before the application terminates.
+
+    **Usage:**
+    Typically used to wrap a standard file handler:
+
+    ```python
+    file_handler = MermaidFileHandler("trace.mmd")
+    async_handler = AsyncMermaidHandler(handlers=[file_handler])
+    logger.addHandler(async_handler)
+    ```
     """
 
     def __init__(self, handlers: List[logging.Handler], queue_size: int = 1000):
         """
-        Initialize the async handler.
+        Initialize the asynchronous handler infrastructure.
+
+        This setup involves three main steps:
+        1.  Creating a bounded queue to hold log records.
+        2.  Initializing the parent QueueHandler.
+        3.  Starting a background QueueListener thread to process the queue.
 
         Args:
-            handlers: A list of handlers that should receive the logs from the queue.
-                      These are typically MermaidFileHandler instances.
-            queue_size: The maximum size of the queue. Default is 1000.
-                       If the queue fills up, new log records may be dropped.
+            handlers (List[logging.Handler]): A list of "real" handlers that will
+                eventually write the logs (e.g., to a file or stream). The background
+                listener will pass records to these handlers.
+            queue_size (int): The maximum number of log records the queue can hold.
+                Defaults to 1000.
+                *Trade-off*: A larger queue consumes more memory but handles larger
+                bursts. A smaller queue saves memory but increases the risk of
+                dropped logs if the consumer falls behind.
         """
-        # Create a bounded queue with the specified size
+        # 1. Create a bounded queue (Producer-Consumer buffer).
+        # We use a bounded queue to prevent uncontrolled memory growth if the
+        # consumer (writer) cannot keep up with the producer (application).
         self._log_queue: queue.Queue[logging.LogRecord] = queue.Queue(queue_size)
         self._queue_size = queue_size
 
-        # Initialize parent QueueHandler with our queue
+        # 2. Initialize the parent QueueHandler.
+        # This configures self.queue, which is used by the emit() method.
         super().__init__(self._log_queue)
 
-        # Initialize QueueListener to process records from the queue
-        # It starts an internal thread to monitor the queue
-        # respect_handler_level=True ensures the target handlers' log levels are respected
+        # 3. Initialize and start the QueueListener (The Consumer).
+        # The QueueListener runs in a separate daemon thread. It continuously:
+        #   a. Blocks waiting for a record from the queue.
+        #   b. Retrieves the record.
+        #   c. Passes it to the provided 'handlers'.
+        #
+        # respect_handler_level=True ensures that if the underlying handler is set
+        # to ERROR but the logger is INFO, the underlying handler won't write INFO logs.
         self._listener: Optional[logging.handlers.QueueListener] = (
             logging.handlers.QueueListener(
                 self._log_queue, *handlers, respect_handler_level=True
             )
         )
 
-        # Start the listener thread
+        # Start the background worker thread.
         self._listener.start()
 
-        # Register stop method to be called on program exit
-        # This ensures all pending logs are written to disk before termination
+        # 4. Ensure Graceful Shutdown.
+        # Register the stop method to be called automatically when the Python
+        # interpreter exits. This is critical for flushing the queue so no logs are lost.
         atexit.register(self.stop)
 
     def emit(self, record: logging.LogRecord) -> None:
         """
-        Emit a log record to the queue with a timeout and drop policy.
+        Emit a log record to the queue (Producer action).
+
+        This method overrides the standard logging emit to implement a non-blocking
+        strategy with a fallback.
 
-        If the queue is full, this method will attempt to put the record with
-        a short timeout. If that fails, it will drop the record and print a warning.
+        **Logic Flow:**
+        1.  Attempt to put the record into the queue.
+        2.  Use a short timeout (0.1s) to avoid blocking the main application
+            indefinitely if the queue is full (backpressure).
+        3.  If the queue remains full after the timeout, drop the record to preserve
+            application stability, but print a warning to stderr.
 
         Args:
-            record: The log record to emit
+            record (logging.LogRecord): The log event to be processed.
         """
-        from typing import cast
-
         try:
-            # Try to put the record in the queue with a short timeout (0.1 seconds)
-            # This prevents the main thread from blocking indefinitely if the queue is full
-            # Use cast to tell Mypy this is a queue.Queue instance
+            # We explicitly cast self.queue because QueueHandler.queue is typed
+            # generically in some stubs, but we know it's a queue.Queue.
             queue_instance = cast(queue.Queue[logging.LogRecord], self.queue)
+
+            # Attempt to enqueue the record.
+            # block=True, timeout=0.1: We wait briefly for a slot to open up.
+            # This balances "trying to save the log" vs "not freezing the app".
             queue_instance.put(record, block=True, timeout=0.1)
+
         except queue.Full:
-            # If queue is full, log a warning and drop the record
+            # **Queue Overflow Handling**
+            # If we reach here, the consumer (writer) is too slow or the burst
+            # is too large. We must drop data to keep the application running.
+
+            # Only warn for important logs to avoid spamming stderr.
             if record.levelno >= logging.WARNING:
-                # Avoid infinite recursion by not using self.logger
+                # We use print() instead of logging to avoid infinite recursion
+                # (logging about a logging failure).
                 print(
-                    f"WARNING: AsyncMermaidHandler queue is full (size: {self._queue_size}), dropping log record: {record.msg}"
+                    f"WARNING: AsyncMermaidHandler queue is full (size: {self._queue_size}), "
+                    f"dropping log record: {record.msg}"
                 )
 
     def stop(self) -> None:
         """
-        Stops the listener and flushes all pending logs from the queue.
+        Clean up resources and flush pending logs.
+
+        This method is called automatically via atexit or can be called manually.
 
-        This method is registered with `atexit` to ensure that all pending logs
-        are written to disk before the application terminates.
+        **Shutdown Sequence:**
+        1.  Check if the listener is active.
+        2.  Call `listener.stop()`. This sends a special "sentinel" (None) to the queue.
+        3.  The background thread sees the sentinel, stops waiting for new logs,
+            processes any remaining items in the queue, and then terminates.
+        4.  Explicitly flush all underlying handlers to ensure stateful formatters
+            write their final buffered events.
         """
         if self._listener:
+            # We keep a reference to handlers to flush them after listener stops
+            handlers = self._listener.handlers
             try:
-                # Stop the listener - this will process all remaining records in the queue
+                # Stop the listener. This blocks until the listener thread joins,
+                # ensuring all records currently in the queue are processed.
                 self._listener.stop()
                 self._listener = None
-            except queue.Full:
-                # Handle case where queue is full when trying to put sentinel value
-                # The listener thread may still be processing, but we can safely exit
+
+                # Crucial step for stateful formatters:
+                # After the listener has finished emitting all records from the queue,
+                # we must tell the handlers to flush their internal buffers.
+                for handler in handlers:
+                    try:
+                        handler.flush()
+                    except Exception:
+                        pass
+            except Exception:
+                # We catch generic exceptions here because during interpreter shutdown,
+                # some modules (like queue) might already be partially unloaded.
                 pass

mermaid_trace/handlers/mermaid_handler.py

@@ -7,22 +7,118 @@ ensures thread-safe file writing.
 """
 
 import logging
+import logging.handlers
 import os
+from typing import Any, Optional, TYPE_CHECKING
 
+if TYPE_CHECKING:
+    pass
 
-class MermaidFileHandler(logging.FileHandler):
+
+class MermaidHandlerMixin:
     """
-    A custom logging handler that writes `FlowEvent` objects to a Mermaid (.mmd) file.
+    Mixin to provide Mermaid-specific logic to logging handlers.
+    """
+
+    title: str
+    terminator: str
+    # These are provided by logging.Handler or its subclasses
+    formatter: Optional[logging.Formatter]
+    stream: Any
 
-    Strategy & Optimization:
-    1. **Inheritance**: Inherits from `logging.FileHandler` to leverage robust,
-       thread-safe file writing capabilities (locking, buffering) provided by the stdlib.
-    2. **Header Management**: Automatically handles the Mermaid file header
-       (`sequenceDiagram`, `title`, `autonumber`) to ensure the output file
-       is a valid Mermaid document. It smartly detects if the file is new or
-       being appended to.
-    3. **Deferred Formatting**: The actual string conversion happens in the `emit`
-       method (via the formatter), keeping the handler focused on I/O.
+    def _write_header(self) -> None:
+        """
+        Writes the initial Mermaid syntax lines to the file.
+        """
+        # Default header if no formatter is available
+        header = f"sequenceDiagram\n    title {self.title}\n    autonumber\n\n"
+
+        if self.formatter and hasattr(self.formatter, "get_header"):
+            try:
+                # Use formatter's header if it provides one
+                header = getattr(self.formatter, "get_header")(self.title)
+                # Ensure it ends with at least one newline for safety
+                if not header.endswith("\n"):
+                    header += "\n"
+            except Exception:
+                # Fallback if formatter fails
+                pass
+
+        if self.stream:
+            self.stream.write(header)
+            # Ensure it's physically on disk before any logs follow
+            self.stream.flush()
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """
+        Process a log record and write it to the Mermaid file.
+        Handles rotation if the parent class supports it.
+        """
+        # Only process records that contain our structured FlowEvent data
+        if not hasattr(record, "flow_event"):
+            return
+
+        try:
+            # 1. Handle Rotation (for RotatingFileHandler and TimedRotatingFileHandler)
+            if hasattr(self, "shouldRollover") and getattr(self, "shouldRollover")(
+                record
+            ):
+                getattr(self, "doRollover")()
+
+            # 2. Ensure stream is open (handles delay=True)
+            if self.stream is None:
+                if hasattr(self, "_open"):
+                    self.stream = getattr(self, "_open")()
+
+            # 3. Check if we need to write the header.
+            # If the file is empty (position 0), it's either a new file,
+            # an empty existing file, or a freshly rotated file.
+            if self.stream and hasattr(self.stream, "tell") and self.stream.tell() == 0:
+                self._write_header()
+
+            # 4. Format the record.
+            # Our custom MermaidFormatter might return an empty string
+            # if it's currently collapsing/buffering repetitive calls.
+            if hasattr(self, "format"):
+                msg = getattr(self, "format")(record)
+                if msg and self.stream:
+                    self.stream.write(msg + self.terminator)
+                    # Note: We do NOT call self.flush() here to allow
+                    # the formatter's collapsing buffer to work correctly.
+        except Exception:
+            if hasattr(self, "handleError"):
+                getattr(self, "handleError")(record)
+
+    def flush(self) -> None:
+        """
+        Flushes both the underlying file stream and any buffered events in the formatter.
+        """
+        if self.formatter and hasattr(self.formatter, "flush"):
+            try:
+                msg = getattr(self.formatter, "flush")()
+                if msg and self.stream:
+                    self.stream.write(msg + self.terminator)
+            except Exception:
+                pass
+
+        # Use hasattr to check if super() has flush, to avoid Mypy errors with mixins
+        super_flush = getattr(super(), "flush", None)
+        if callable(super_flush):
+            super_flush()
+
+    def close(self) -> None:
+        """
+        Ensures all buffered events are written before closing the file.
+        """
+        self.flush()
+        super_close = getattr(super(), "close", None)
+        if callable(super_close):
+            super_close()
+
+
+class MermaidFileHandler(MermaidHandlerMixin, logging.FileHandler):
+    """
+    A custom logging handler that writes `FlowEvent` objects to a Mermaid (.mmd) file.
     """
 
     def __init__(
@@ -33,87 +129,57 @@ class MermaidFileHandler(logging.FileHandler):
         encoding: str = "utf-8",
         delay: bool = False,
     ):
-        """
-        Initialize the Mermaid file handler.
-
-        Args:
-            filename (str): The path to the output .mmd file.
-            title (str, optional): The title of the Mermaid diagram. Defaults to "Log Flow".
-            mode (str, optional): File open mode. 'w' (overwrite) or 'a' (append). Defaults to "a".
-            encoding (str, optional): File encoding. Defaults to "utf-8".
-            delay (bool, optional): If True, file opening is deferred until the first call to emit.
-                                   Useful to avoid creating empty files if no logs occur. Defaults to False.
-        """
-        # Ensure the directory exists to prevent FileNotFoundError when opening the file
         os.makedirs(os.path.dirname(os.path.abspath(filename)) or ".", exist_ok=True)
-
-        # Determine if we need to write a header
-        # Header is written only if:
-        # 1. We are overwriting the file (mode='w'), or
-        # 2. We are appending (mode='a') but the file doesn't exist or is empty
-        should_write_header = False
-        if mode == "w":
-            should_write_header = True
-        elif mode == "a":
-            if not os.path.exists(filename) or os.path.getsize(filename) == 0:
-                should_write_header = True
-
-        # Initialize the parent FileHandler (opens the file unless delay=True)
         super().__init__(filename, mode, encoding, delay)
         self.title = title
+        self.terminator = "\n"
 
-        # Write the header immediately if needed
-        if should_write_header:
-            self._write_header()
-
-    def _write_header(self) -> None:
-        """
-        Writes the initial Mermaid syntax lines to the file.
 
-        This setup is required for Mermaid JS or Live Editor to render the diagram correctly.
-        It defines:
-        - Diagram type (sequenceDiagram)
-        - Title of the diagram
-        - Autonumbering of steps
-
-        Thread Safety: Uses the handler's internal lock to prevent concurrent writes
-        when delay=True, ensuring the header is written only once.
-        """
-        # Use the handler's internal lock to ensure thread safety
-        assert self.lock is not None, "Handler lock should always be initialized"
-        with self.lock:
-            # Write to the existing stream if available, otherwise open temporarily
-            if self.stream:
-                # Stream is already open (delay=False or emit() has been called)
-                self.stream.write("sequenceDiagram\n")
-                self.stream.write(f"    title {self.title}\n")
-                self.stream.write("    autonumber\n")
-                # Flush ensures the header is written to disk immediately
-                self.flush()
-            else:
-                # Handle delay=True case: file not yet opened
-                # Temporarily open the file just to write the header
-                with open(self.baseFilename, self.mode, encoding=self.encoding) as f:
-                    f.write("sequenceDiagram\n")
-                    f.write(f"    title {self.title}\n")
-                    f.write("    autonumber\n")
+class RotatingMermaidFileHandler(
+    MermaidHandlerMixin, logging.handlers.RotatingFileHandler
+):
+    """
+    Rotating version of the MermaidFileHandler.
+    """
 
-    def emit(self, record: logging.LogRecord) -> None:
-        """
-        Process a log record and write it to the Mermaid file.
+    def __init__(
+        self,
+        filename: str,
+        title: str = "Log Flow",
+        mode: str = "a",
+        maxBytes: int = 0,
+        backupCount: int = 0,
+        encoding: str = "utf-8",
+        delay: bool = False,
+    ):  # noqa: PLR0913
+        os.makedirs(os.path.dirname(os.path.abspath(filename)) or ".", exist_ok=True)
+        super().__init__(filename, mode, maxBytes, backupCount, encoding, delay)
+        self.title = title
+        self.terminator = "\n"
 
-        This method overrides the parent's emit method to filter out non-FlowEvent records.
 
-        Optimization:
-        - Checks for `flow_event` attribute first, acting as a high-performance filter
-        - Only processes records containing structured FlowEvent data
-        - Delegates actual writing to parent's emit() method, which handles locking and flushing
+class TimedRotatingMermaidFileHandler(
+    MermaidHandlerMixin, logging.handlers.TimedRotatingFileHandler
+):
+    """
+    Timed rotating version of the MermaidFileHandler.
+    """
 
-        Args:
-            record (logging.LogRecord): The log record to process
-        """
-        # Only process records that contain our structured FlowEvent data
-        # This allows the handler to be attached to the root logger without processing
-        # irrelevant system logs
-        if hasattr(record, "flow_event"):
-            super().emit(record)
+    def __init__(
+        self,
+        filename: str,
+        title: str = "Log Flow",
+        when: str = "h",
+        interval: int = 1,
+        backupCount: int = 0,
+        encoding: str = "utf-8",
+        delay: bool = False,
+        utc: bool = False,
+        atTime: Any = None,
+    ):  # noqa: PLR0913
+        os.makedirs(os.path.dirname(os.path.abspath(filename)) or ".", exist_ok=True)
+        super().__init__(
+            filename, when, interval, backupCount, encoding, delay, utc, atTime
+        )
+        self.title = title
+        self.terminator = "\n"

mermaid_trace/integrations/__init__.py

@@ -0,0 +1,4 @@
+"""
+MermaidTrace integrations package.
+Contains middleware and adapters for third-party frameworks.
+"""