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

mermaid_trace/core/context.py

@@ -1,3 +1,12 @@
+"""
+Log Context Management Module
+
+This module provides a thread-safe, async-friendly context management system
+for tracking execution context across the application. It uses Python's ContextVar
+mechanism to ensure proper context propagation in both synchronous and asynchronous
+environments.
+"""
+
 from contextvars import ContextVar, Token
 from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncIterator, Dict, Iterator
@@ -9,16 +18,14 @@ 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
-    correct context propagation in asynchronous (asyncio) environments.
-    Unlike `threading.local()`, `ContextVar` works natively with Python's async/await
+    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
     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.
-    # The default value is implicitly an empty state if not set (handled in _get_store).
+    # ContextVar stores a dictionary unique to the current execution context (Task/Thread)
+    # The name "log_context" is used for debugging purposes
     _context_store: ContextVar[Dict[str, Any]] = ContextVar("log_context")
 
     @classmethod
@@ -27,13 +34,19 @@ class LogContext:
         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
-        and ensures there's always a valid dictionary to work with.
+        it creates a fresh empty dictionary, sets it to the contextvar,
+        and returns it. This prevents LookupError and ensures there's
+        always a valid dictionary to work with.
+
+        Returns:
+            Dict[str, Any]: Current context dictionary for the execution flow
         """
         try:
             return cls._context_store.get()
         except LookupError:
-            return {}
+            empty_dict: Dict[str, Any] = {}
+            cls._context_store.set(empty_dict)
+            return empty_dict
 
     @classmethod
     def set(cls, key: str, value: Any) -> None:
@@ -42,10 +55,14 @@ class LogContext:
 
         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).
-        3. Update the copy.
-        4. Re-set the ContextVar with the new dictionary.
+        1. Retrieve the current dictionary using _get_store()
+        2. Create a shallow copy to avoid affecting parent contexts
+        3. Update the copy with the new key-value pair
+        4. Re-set the ContextVar with the new dictionary
+
+        Args:
+            key (str): Name of the context variable to set
+            value (Any): Value to associate with the key
         """
         ctx = cls._get_store().copy()
         ctx[key] = value
@@ -57,7 +74,10 @@ class LogContext:
         Updates multiple keys in the current context at once.
 
         This follows the same Copy-Update-Set pattern as `set()` to maintain
-        context isolation.
+        context isolation between different execution flows.
+
+        Args:
+            data (Dict[str, Any]): Dictionary of key-value pairs to update in context
         """
         if not data:
             return
@@ -69,6 +89,13 @@ class LogContext:
     def get(cls, key: str, default: Any = None) -> Any:
         """
         Retrieves a value from the current context safely.
+
+        Args:
+            key (str): Name of the context variable to retrieve
+            default (Any, optional): Default value if key doesn't exist. Defaults to None.
+
+        Returns:
+            Any: Value associated with the key, or default if key not found
         """
         return cls._get_store().get(key, default)
 
@@ -76,6 +103,9 @@ class LogContext:
     def get_all(cls) -> Dict[str, Any]:
         """
         Returns a copy of the entire context dictionary.
+
+        Returns:
+            Dict[str, Any]: Complete copy of the current context
         """
         return cls._get_store().copy()
 
@@ -92,11 +122,17 @@ class LogContext:
             # 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.
+            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
-               before the block entered.
+               before the block entered
+
+        Args:
+            data (Dict[str, Any]): Dictionary of context values to set within the scope
+
+        Yields:
+            None: Control to the block using this context manager
         """
         current_ctx = cls._get_store().copy()
         current_ctx.update(data)
@@ -104,7 +140,7 @@ class LogContext:
         try:
             yield
         finally:
-            # Crucial: Reset restores the context to what it was before .set()
+            # Reset restores context to state before .set() was called
             cls._context_store.reset(token)
 
     @classmethod
@@ -120,6 +156,12 @@ class LogContext:
         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.
+
+        Args:
+            data (Dict[str, Any]): Dictionary of context values to set within the scope
+
+        Yields:
+            None: Control to the async block using this context manager
         """
         current_ctx = cls._get_store().copy()
         current_ctx.update(data)
@@ -137,6 +179,12 @@ class LogContext:
         """
         Replaces the entire context with the provided data.
         Returns a Token that can be used to manually reset the context later.
+
+        Args:
+            data (Dict[str, Any]): New context dictionary to replace the current one
+
+        Returns:
+            Token[Dict[str, Any]]: Token for resetting context to previous state
         """
         return cls._context_store.set(data.copy())
 
@@ -144,21 +192,30 @@ class LogContext:
     def reset(cls, token: Token[Dict[str, Any]]) -> None:
         """
         Manually resets the context using a Token obtained from `set` or `set_all`.
+
+        Args:
+            token (Token[Dict[str, Any]]): Token returned by set_all() method
         """
         cls._context_store.reset(token)
 
     @classmethod
     def current_participant(cls) -> str:
         """
-        Helper to get the 'participant' field, representing the current active object/module.
+        Helper method to get the 'participant' field, representing the current active object/module.
         Defaults to 'Unknown' if not set.
+
+        Returns:
+            str: Name of the current participant in the trace flow
         """
         return str(cls.get("participant", "Unknown"))
 
     @classmethod
     def set_participant(cls, name: str) -> None:
         """
-        Helper to set the 'participant' field.
+        Helper method to set the 'participant' field.
+
+        Args:
+            name (str): Name of the participant to set
         """
         cls.set("participant", name)
 
@@ -170,9 +227,12 @@ class LogContext:
         Lazy Initialization Logic:
         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.
+        1. A trace ID is always available when asked for
         2. Once generated, the same ID persists for the duration of the context
-           (unless manually changed), linking all subsequent logs together.
+           (unless manually changed), linking all subsequent logs together
+
+        Returns:
+            str: Unique trace ID for the current execution flow
         """
         tid = cls.get("trace_id")
         if not tid: