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

mermaid_trace/core/context.py

@@ -1,50 +1,68 @@
+"""
+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
 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 
-    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 
+
+    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
     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
     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 
-        and ensures there's always a valid dictionary to work with.
+
+        If the context variable has not been set in the current context,
+        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:
         """
         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).
-        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
@@ -54,9 +72,12 @@ 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 
-        context isolation.
+
+        This follows the same Copy-Update-Set pattern as `set()` to maintain
+        context isolation between different execution flows.
+
+        Args:
+            data (Dict[str, Any]): Dictionary of key-value pairs to update in context
         """
         if not data:
             return
@@ -68,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)
 
@@ -75,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()
 
@@ -83,19 +114,25 @@ 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 
-               before the block entered.
+            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
+
+        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)
@@ -103,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
@@ -111,14 +148,20 @@ 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.
+
+        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)
@@ -136,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())
 
@@ -143,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)
 
@@ -165,13 +223,16 @@ 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 
-           (unless manually changed), linking all subsequent logs together.
+        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
+
+        Returns:
+            str: Unique trace ID for the current execution flow
         """
         tid = cls.get("trace_id")
         if not tid: