mermaid-trace 0.3.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mermaid_trace/__init__.py +70 -0
- mermaid_trace/cli.py +193 -0
- mermaid_trace/core/__init__.py +0 -0
- mermaid_trace/core/context.py +187 -0
- mermaid_trace/core/decorators.py +252 -0
- mermaid_trace/core/events.py +75 -0
- mermaid_trace/core/formatter.py +72 -0
- mermaid_trace/handlers/mermaid_handler.py +86 -0
- mermaid_trace/integrations/fastapi.py +122 -0
- mermaid_trace/py.typed +0 -0
- mermaid_trace-0.3.1.dist-info/METADATA +163 -0
- mermaid_trace-0.3.1.dist-info/RECORD +15 -0
- mermaid_trace-0.3.1.dist-info/WHEEL +4 -0
- mermaid_trace-0.3.1.dist-info/entry_points.txt +2 -0
- mermaid_trace-0.3.1.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""
|
|
2
|
+
MermaidTrace: Visualize your Python code execution flow as Mermaid Sequence Diagrams.
|
|
3
|
+
|
|
4
|
+
This package provides tools to automatically trace function calls and generate
|
|
5
|
+
Mermaid-compatible sequence diagrams (.mmd files). It is designed to help
|
|
6
|
+
developers understand the flow of their applications, debug complex interactions,
|
|
7
|
+
and document system behavior.
|
|
8
|
+
|
|
9
|
+
Key Components:
|
|
10
|
+
- `trace`: A decorator to instrument functions for tracing.
|
|
11
|
+
- `LogContext`: Manages execution context (like thread-local storage) to track
|
|
12
|
+
caller/callee relationships across async tasks and threads.
|
|
13
|
+
- `configure_flow`: Sets up the logging handler to write diagrams to a file.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from .core.decorators import trace_interaction, trace
|
|
17
|
+
from .handlers.mermaid_handler import MermaidFileHandler
|
|
18
|
+
from .core.events import FlowEvent
|
|
19
|
+
from .core.context import LogContext
|
|
20
|
+
from .core.formatter import MermaidFormatter
|
|
21
|
+
# We don't import integrations by default to avoid hard dependencies
|
|
22
|
+
# Integrations (like FastAPI) must be imported explicitly by the user if needed.
|
|
23
|
+
|
|
24
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
25
|
+
|
|
26
|
+
import logging
|
|
27
|
+
|
|
28
|
+
def configure_flow(output_file: str = "flow.mmd") -> logging.Logger:
|
|
29
|
+
"""
|
|
30
|
+
Configures the flow logger to output to a Mermaid file.
|
|
31
|
+
|
|
32
|
+
This function sets up the logging infrastructure required to capture
|
|
33
|
+
trace events and write them to the specified output file. It should
|
|
34
|
+
be called once at the start of your application.
|
|
35
|
+
|
|
36
|
+
Args:
|
|
37
|
+
output_file (str): The absolute or relative path to the output .mmd file.
|
|
38
|
+
Defaults to "flow.mmd" in the current directory.
|
|
39
|
+
|
|
40
|
+
Returns:
|
|
41
|
+
logging.Logger: The configured logger instance used for flow tracing.
|
|
42
|
+
"""
|
|
43
|
+
# Get the specific logger used by the tracing decorators
|
|
44
|
+
logger = logging.getLogger("mermaid_trace.flow")
|
|
45
|
+
logger.setLevel(logging.INFO)
|
|
46
|
+
|
|
47
|
+
# Remove existing handlers to avoid duplicate logs if configured multiple times
|
|
48
|
+
if logger.hasHandlers():
|
|
49
|
+
logger.handlers.clear()
|
|
50
|
+
|
|
51
|
+
# Create and attach the custom handler that writes Mermaid syntax
|
|
52
|
+
handler = MermaidFileHandler(output_file)
|
|
53
|
+
|
|
54
|
+
# Use the custom formatter to convert FlowEvents to Mermaid strings
|
|
55
|
+
handler.setFormatter(MermaidFormatter())
|
|
56
|
+
|
|
57
|
+
logger.addHandler(handler)
|
|
58
|
+
|
|
59
|
+
return logger
|
|
60
|
+
|
|
61
|
+
try:
|
|
62
|
+
# Attempt to retrieve the installed package version
|
|
63
|
+
__version__ = version("mermaid-trace")
|
|
64
|
+
except PackageNotFoundError:
|
|
65
|
+
# Fallback version if the package is not installed (e.g., local development)
|
|
66
|
+
__version__ = "0.0.0"
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
# Export public API for easy access
|
|
70
|
+
__all__ = ["trace_interaction", "trace", "configure_flow", "MermaidFileHandler", "LogContext", "FlowEvent", "MermaidFormatter"]
|
mermaid_trace/cli.py
ADDED
|
@@ -0,0 +1,193 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import http.server
|
|
3
|
+
import socketserver
|
|
4
|
+
import webbrowser
|
|
5
|
+
import sys
|
|
6
|
+
import os
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import Type, Any
|
|
9
|
+
|
|
10
|
+
try:
|
|
11
|
+
from watchdog.observers import Observer
|
|
12
|
+
from watchdog.events import FileSystemEventHandler
|
|
13
|
+
HAS_WATCHDOG = True
|
|
14
|
+
except ImportError:
|
|
15
|
+
HAS_WATCHDOG = False
|
|
16
|
+
|
|
17
|
+
# HTML Template for the preview page
|
|
18
|
+
HTML_TEMPLATE = """
|
|
19
|
+
<!DOCTYPE html>
|
|
20
|
+
<html lang="en">
|
|
21
|
+
<head>
|
|
22
|
+
<meta charset="UTF-8">
|
|
23
|
+
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
|
24
|
+
<title>MermaidTrace Flow Preview</title>
|
|
25
|
+
<!-- Load Mermaid.js from CDN -->
|
|
26
|
+
<script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
|
|
27
|
+
<style>
|
|
28
|
+
/* Basic styling for readability and layout */
|
|
29
|
+
body {{ font-family: sans-serif; padding: 20px; background: #f4f4f4; }}
|
|
30
|
+
.container {{ background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }}
|
|
31
|
+
h1 {{ color: #333; }}
|
|
32
|
+
#diagram {{ overflow-x: auto; }} /* Allow horizontal scrolling for wide diagrams */
|
|
33
|
+
.refresh-btn {{
|
|
34
|
+
position: fixed; bottom: 20px; right: 20px;
|
|
35
|
+
padding: 10px 20px; background: #007bff; color: white;
|
|
36
|
+
border: none; border-radius: 5px; cursor: pointer; font-size: 16px;
|
|
37
|
+
}}
|
|
38
|
+
.refresh-btn:hover {{ background: #0056b3; }}
|
|
39
|
+
.status {{
|
|
40
|
+
position: fixed; bottom: 20px; left: 20px;
|
|
41
|
+
font-size: 12px; color: #666;
|
|
42
|
+
}}
|
|
43
|
+
</style>
|
|
44
|
+
</head>
|
|
45
|
+
<body>
|
|
46
|
+
<div class="container">
|
|
47
|
+
<h1>MermaidTrace Flow Preview: {filename}</h1>
|
|
48
|
+
<!-- The Mermaid diagram content will be injected here -->
|
|
49
|
+
<div class="mermaid" id="diagram">
|
|
50
|
+
{content}
|
|
51
|
+
</div>
|
|
52
|
+
</div>
|
|
53
|
+
<!-- Button to manually reload the page/diagram -->
|
|
54
|
+
<button class="refresh-btn" onclick="location.reload()">Refresh Diagram</button>
|
|
55
|
+
<div class="status" id="status">Monitoring for changes...</div>
|
|
56
|
+
|
|
57
|
+
<script>
|
|
58
|
+
// Initialize Mermaid with default settings
|
|
59
|
+
mermaid.initialize({{ startOnLoad: true, theme: 'default' }});
|
|
60
|
+
|
|
61
|
+
// Live Reload Logic
|
|
62
|
+
const currentMtime = "{mtime}";
|
|
63
|
+
|
|
64
|
+
function checkUpdate() {{
|
|
65
|
+
fetch('/_status')
|
|
66
|
+
.then(response => response.text())
|
|
67
|
+
.then(mtime => {{
|
|
68
|
+
if (mtime && mtime !== currentMtime) {{
|
|
69
|
+
console.log("File changed, reloading...");
|
|
70
|
+
location.reload();
|
|
71
|
+
}}
|
|
72
|
+
}})
|
|
73
|
+
.catch(err => console.error("Error checking status:", err));
|
|
74
|
+
}}
|
|
75
|
+
|
|
76
|
+
// Poll every 1 second
|
|
77
|
+
setInterval(checkUpdate, 1000);
|
|
78
|
+
</script>
|
|
79
|
+
</body>
|
|
80
|
+
</html>
|
|
81
|
+
"""
|
|
82
|
+
|
|
83
|
+
def _create_handler(filename: str, path: Path) -> Type[http.server.SimpleHTTPRequestHandler]:
|
|
84
|
+
"""Helper to create the request handler class with closure over filename/path."""
|
|
85
|
+
class Handler(http.server.SimpleHTTPRequestHandler):
|
|
86
|
+
"""
|
|
87
|
+
Custom Request Handler to serve the generated HTML dynamically.
|
|
88
|
+
"""
|
|
89
|
+
def log_message(self, format: str, *args: Any) -> None:
|
|
90
|
+
# Suppress default logging to keep console clean
|
|
91
|
+
pass
|
|
92
|
+
|
|
93
|
+
def do_GET(self) -> None:
|
|
94
|
+
"""
|
|
95
|
+
Handle GET requests.
|
|
96
|
+
Serves the HTML wrapper for the root path ('/').
|
|
97
|
+
"""
|
|
98
|
+
if self.path == "/":
|
|
99
|
+
self.send_response(200)
|
|
100
|
+
self.send_header("Content-type", "text/html")
|
|
101
|
+
self.end_headers()
|
|
102
|
+
|
|
103
|
+
try:
|
|
104
|
+
# Read the current content of the mermaid file
|
|
105
|
+
content = path.read_text(encoding="utf-8")
|
|
106
|
+
mtime = str(path.stat().st_mtime)
|
|
107
|
+
except Exception as e:
|
|
108
|
+
# Fallback if reading fails (e.g., file locked)
|
|
109
|
+
content = f"sequenceDiagram\nNote right of Error: Failed to read file: {e}"
|
|
110
|
+
mtime = "0"
|
|
111
|
+
|
|
112
|
+
# Inject content into the HTML template
|
|
113
|
+
html = HTML_TEMPLATE.format(filename=filename, content=content, mtime=mtime)
|
|
114
|
+
self.wfile.write(html.encode("utf-8"))
|
|
115
|
+
|
|
116
|
+
elif self.path == "/_status":
|
|
117
|
+
# API endpoint for client to check file status
|
|
118
|
+
self.send_response(200)
|
|
119
|
+
self.send_header("Content-type", "text/plain")
|
|
120
|
+
self.end_headers()
|
|
121
|
+
try:
|
|
122
|
+
mtime = str(path.stat().st_mtime)
|
|
123
|
+
except OSError:
|
|
124
|
+
mtime = "0"
|
|
125
|
+
self.wfile.write(mtime.encode("utf-8"))
|
|
126
|
+
|
|
127
|
+
else:
|
|
128
|
+
# Serve static files if needed, or return 404
|
|
129
|
+
super().do_GET()
|
|
130
|
+
return Handler
|
|
131
|
+
|
|
132
|
+
def serve(filename: str, port: int = 8000) -> None:
|
|
133
|
+
"""
|
|
134
|
+
Starts a local HTTP server to preview the Mermaid diagram.
|
|
135
|
+
"""
|
|
136
|
+
path = Path(filename)
|
|
137
|
+
if not path.exists():
|
|
138
|
+
print(f"Error: File '{filename}' not found.")
|
|
139
|
+
sys.exit(1)
|
|
140
|
+
|
|
141
|
+
# Setup Watchdog if available
|
|
142
|
+
observer = None
|
|
143
|
+
if HAS_WATCHDOG:
|
|
144
|
+
class FileChangeHandler(FileSystemEventHandler):
|
|
145
|
+
def on_modified(self, event: Any) -> None:
|
|
146
|
+
if not event.is_directory and os.path.abspath(event.src_path) == str(path.resolve()):
|
|
147
|
+
print(f"[Watchdog] File changed: {filename}")
|
|
148
|
+
|
|
149
|
+
print("Initializing file watcher...")
|
|
150
|
+
observer = Observer()
|
|
151
|
+
observer.schedule(FileChangeHandler(), path=str(path.parent), recursive=False)
|
|
152
|
+
observer.start()
|
|
153
|
+
else:
|
|
154
|
+
print("Watchdog not installed. Falling back to polling mode (client-side only).")
|
|
155
|
+
|
|
156
|
+
HandlerClass = _create_handler(filename, path)
|
|
157
|
+
|
|
158
|
+
print(f"Serving {filename} at http://localhost:{port}")
|
|
159
|
+
print("Press Ctrl+C to stop.")
|
|
160
|
+
|
|
161
|
+
# Open browser automatically to the server URL
|
|
162
|
+
webbrowser.open(f"http://localhost:{port}")
|
|
163
|
+
|
|
164
|
+
# Start the TCP server
|
|
165
|
+
with socketserver.ThreadingTCPServer(("", port), HandlerClass) as httpd:
|
|
166
|
+
try:
|
|
167
|
+
httpd.serve_forever()
|
|
168
|
+
except KeyboardInterrupt:
|
|
169
|
+
print("\nStopping server...")
|
|
170
|
+
if observer:
|
|
171
|
+
observer.stop()
|
|
172
|
+
observer.join()
|
|
173
|
+
httpd.server_close()
|
|
174
|
+
|
|
175
|
+
def main() -> None:
|
|
176
|
+
"""
|
|
177
|
+
Entry point for the CLI application.
|
|
178
|
+
"""
|
|
179
|
+
parser = argparse.ArgumentParser(description="MermaidTrace CLI")
|
|
180
|
+
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
181
|
+
|
|
182
|
+
# 'serve' command definition
|
|
183
|
+
serve_parser = subparsers.add_parser("serve", help="Serve a Mermaid file in the browser")
|
|
184
|
+
serve_parser.add_argument("file", help="Path to the .mmd file")
|
|
185
|
+
serve_parser.add_argument("--port", type=int, default=8000, help="Port to bind to (default: 8000)")
|
|
186
|
+
|
|
187
|
+
args = parser.parse_args()
|
|
188
|
+
|
|
189
|
+
if args.command == "serve":
|
|
190
|
+
serve(args.file, args.port)
|
|
191
|
+
|
|
192
|
+
if __name__ == "__main__":
|
|
193
|
+
main()
|
|
File without changes
|
|
@@ -0,0 +1,187 @@
|
|
|
1
|
+
from contextvars import ContextVar, Token
|
|
2
|
+
from contextlib import asynccontextmanager, contextmanager
|
|
3
|
+
from typing import Any, AsyncIterator, Dict, Iterator
|
|
4
|
+
import uuid
|
|
5
|
+
|
|
6
|
+
class LogContext:
|
|
7
|
+
"""
|
|
8
|
+
Manages global context information for logging (e.g., request_id, user_id, current_participant).
|
|
9
|
+
|
|
10
|
+
This class utilizes `contextvars.ContextVar` to ensure thread-safety and
|
|
11
|
+
correct context propagation in asynchronous (asyncio) environments.
|
|
12
|
+
Unlike `threading.local()`, `ContextVar` works natively with Python's async/await
|
|
13
|
+
event loop, ensuring that context is preserved across `await` points but isolated
|
|
14
|
+
between different concurrent tasks.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
# ContextVar is the key mechanism here.
|
|
18
|
+
# It stores a dictionary unique to the current execution context (Task/Thread).
|
|
19
|
+
# "log_context" is the name of the variable, useful for debugging.
|
|
20
|
+
# The default value is implicitly an empty state if not set (handled in _get_store).
|
|
21
|
+
_context_store: ContextVar[Dict[str, Any]] = ContextVar("log_context")
|
|
22
|
+
|
|
23
|
+
@classmethod
|
|
24
|
+
def _get_store(cls) -> Dict[str, Any]:
|
|
25
|
+
"""
|
|
26
|
+
Retrieves the current context dictionary.
|
|
27
|
+
|
|
28
|
+
If the context variable has not been set in the current context,
|
|
29
|
+
it returns a fresh empty dictionary. This prevents LookupError
|
|
30
|
+
and ensures there's always a valid dictionary to work with.
|
|
31
|
+
"""
|
|
32
|
+
try:
|
|
33
|
+
return cls._context_store.get()
|
|
34
|
+
except LookupError:
|
|
35
|
+
return {}
|
|
36
|
+
|
|
37
|
+
@classmethod
|
|
38
|
+
def set(cls, key: str, value: Any) -> None:
|
|
39
|
+
"""
|
|
40
|
+
Sets a specific key-value pair in the current context.
|
|
41
|
+
|
|
42
|
+
Important: ContextVars are immutable collections. To modify the context,
|
|
43
|
+
we must:
|
|
44
|
+
1. Retrieve the current dictionary.
|
|
45
|
+
2. Create a shallow copy (to avoid affecting parent contexts if we were reusing the object).
|
|
46
|
+
3. Update the copy.
|
|
47
|
+
4. Re-set the ContextVar with the new dictionary.
|
|
48
|
+
"""
|
|
49
|
+
ctx = cls._get_store().copy()
|
|
50
|
+
ctx[key] = value
|
|
51
|
+
cls._context_store.set(ctx)
|
|
52
|
+
|
|
53
|
+
@classmethod
|
|
54
|
+
def update(cls, data: Dict[str, Any]) -> None:
|
|
55
|
+
"""
|
|
56
|
+
Updates multiple keys in the current context at once.
|
|
57
|
+
|
|
58
|
+
This follows the same Copy-Update-Set pattern as `set()` to maintain
|
|
59
|
+
context isolation.
|
|
60
|
+
"""
|
|
61
|
+
if not data:
|
|
62
|
+
return
|
|
63
|
+
ctx = cls._get_store().copy()
|
|
64
|
+
ctx.update(data)
|
|
65
|
+
cls._context_store.set(ctx)
|
|
66
|
+
|
|
67
|
+
@classmethod
|
|
68
|
+
def get(cls, key: str, default: Any = None) -> Any:
|
|
69
|
+
"""
|
|
70
|
+
Retrieves a value from the current context safely.
|
|
71
|
+
"""
|
|
72
|
+
return cls._get_store().get(key, default)
|
|
73
|
+
|
|
74
|
+
@classmethod
|
|
75
|
+
def get_all(cls) -> Dict[str, Any]:
|
|
76
|
+
"""
|
|
77
|
+
Returns a copy of the entire context dictionary.
|
|
78
|
+
"""
|
|
79
|
+
return cls._get_store().copy()
|
|
80
|
+
|
|
81
|
+
@classmethod
|
|
82
|
+
@contextmanager
|
|
83
|
+
def scope(cls, data: Dict[str, Any]) -> Iterator[None]:
|
|
84
|
+
"""
|
|
85
|
+
Synchronous context manager for temporary context updates.
|
|
86
|
+
|
|
87
|
+
Usage:
|
|
88
|
+
with LogContext.scope({"user_id": 123}):
|
|
89
|
+
# user_id is 123 here
|
|
90
|
+
some_function()
|
|
91
|
+
# user_id reverts to previous value (or disappears) here
|
|
92
|
+
|
|
93
|
+
Mechanism:
|
|
94
|
+
1. Copies current context and updates it with new data.
|
|
95
|
+
2. Sets the ContextVar to this new state, receiving a `Token`.
|
|
96
|
+
3. Yields control to the block.
|
|
97
|
+
4. Finally, uses the `Token` to reset the ContextVar to its exact state
|
|
98
|
+
before the block entered.
|
|
99
|
+
"""
|
|
100
|
+
current_ctx = cls._get_store().copy()
|
|
101
|
+
current_ctx.update(data)
|
|
102
|
+
token = cls._context_store.set(current_ctx)
|
|
103
|
+
try:
|
|
104
|
+
yield
|
|
105
|
+
finally:
|
|
106
|
+
# Crucial: Reset restores the context to what it was before .set()
|
|
107
|
+
cls._context_store.reset(token)
|
|
108
|
+
|
|
109
|
+
@classmethod
|
|
110
|
+
@asynccontextmanager
|
|
111
|
+
async def ascope(cls, data: Dict[str, Any]) -> AsyncIterator[None]:
|
|
112
|
+
"""
|
|
113
|
+
Async context manager for temporary context updates in coroutines.
|
|
114
|
+
|
|
115
|
+
Usage:
|
|
116
|
+
async with LogContext.ascope({"request_id": "abc"}):
|
|
117
|
+
await some_async_function()
|
|
118
|
+
|
|
119
|
+
This is functionally identical to `scope` but designed for `async with` blocks.
|
|
120
|
+
It ensures that even if the code inside `yield` suspends execution (await),
|
|
121
|
+
the context remains valid for that task.
|
|
122
|
+
"""
|
|
123
|
+
current_ctx = cls._get_store().copy()
|
|
124
|
+
current_ctx.update(data)
|
|
125
|
+
token = cls._context_store.set(current_ctx)
|
|
126
|
+
try:
|
|
127
|
+
yield
|
|
128
|
+
finally:
|
|
129
|
+
cls._context_store.reset(token)
|
|
130
|
+
|
|
131
|
+
# Alias for backward compatibility if needed
|
|
132
|
+
ascope_async = ascope
|
|
133
|
+
|
|
134
|
+
@classmethod
|
|
135
|
+
def set_all(cls, data: Dict[str, Any]) -> Token[Dict[str, Any]]:
|
|
136
|
+
"""
|
|
137
|
+
Replaces the entire context with the provided data.
|
|
138
|
+
Returns a Token that can be used to manually reset the context later.
|
|
139
|
+
"""
|
|
140
|
+
return cls._context_store.set(data.copy())
|
|
141
|
+
|
|
142
|
+
@classmethod
|
|
143
|
+
def reset(cls, token: Token[Dict[str, Any]]) -> None:
|
|
144
|
+
"""
|
|
145
|
+
Manually resets the context using a Token obtained from `set` or `set_all`.
|
|
146
|
+
"""
|
|
147
|
+
cls._context_store.reset(token)
|
|
148
|
+
|
|
149
|
+
@classmethod
|
|
150
|
+
def current_participant(cls) -> str:
|
|
151
|
+
"""
|
|
152
|
+
Helper to get the 'participant' field, representing the current active object/module.
|
|
153
|
+
Defaults to 'Unknown' if not set.
|
|
154
|
+
"""
|
|
155
|
+
return str(cls.get("participant", "Unknown"))
|
|
156
|
+
|
|
157
|
+
@classmethod
|
|
158
|
+
def set_participant(cls, name: str) -> None:
|
|
159
|
+
"""
|
|
160
|
+
Helper to set the 'participant' field.
|
|
161
|
+
"""
|
|
162
|
+
cls.set("participant", name)
|
|
163
|
+
|
|
164
|
+
@classmethod
|
|
165
|
+
def current_trace_id(cls) -> str:
|
|
166
|
+
"""
|
|
167
|
+
Retrieves the current trace ID for correlating events in a single flow.
|
|
168
|
+
|
|
169
|
+
Lazy Initialization Logic:
|
|
170
|
+
If no trace_id exists in the current context, it generates a new UUIDv4
|
|
171
|
+
and sets it immediately. This ensures that:
|
|
172
|
+
1. A trace ID is always available when asked for.
|
|
173
|
+
2. Once generated, the same ID persists for the duration of the context
|
|
174
|
+
(unless manually changed), linking all subsequent logs together.
|
|
175
|
+
"""
|
|
176
|
+
tid = cls.get("trace_id")
|
|
177
|
+
if not tid:
|
|
178
|
+
tid = str(uuid.uuid4())
|
|
179
|
+
cls.set("trace_id", tid)
|
|
180
|
+
return str(tid)
|
|
181
|
+
|
|
182
|
+
@classmethod
|
|
183
|
+
def set_trace_id(cls, trace_id: str) -> None:
|
|
184
|
+
"""
|
|
185
|
+
Manually sets the trace ID (e.g., from an incoming HTTP request header).
|
|
186
|
+
"""
|
|
187
|
+
cls.set("trace_id", trace_id)
|
|
@@ -0,0 +1,252 @@
|
|
|
1
|
+
import functools
|
|
2
|
+
import logging
|
|
3
|
+
import inspect
|
|
4
|
+
import reprlib
|
|
5
|
+
from typing import Optional, Any, Callable, Tuple, Dict, Union, TypeVar, cast, overload
|
|
6
|
+
|
|
7
|
+
from .events import FlowEvent
|
|
8
|
+
from .context import LogContext
|
|
9
|
+
|
|
10
|
+
FLOW_LOGGER_NAME = "mermaid_trace.flow"
|
|
11
|
+
|
|
12
|
+
# Define generic type variable for the decorated function
|
|
13
|
+
F = TypeVar("F", bound=Callable[..., Any])
|
|
14
|
+
|
|
15
|
+
def get_flow_logger() -> logging.Logger:
|
|
16
|
+
"""Returns the dedicated logger for flow events."""
|
|
17
|
+
return logging.getLogger(FLOW_LOGGER_NAME)
|
|
18
|
+
|
|
19
|
+
def _safe_repr(obj: Any, max_len: int = 50) -> str:
|
|
20
|
+
"""
|
|
21
|
+
Safely creates a string representation of an object.
|
|
22
|
+
|
|
23
|
+
Prevents massive log files by truncating long strings/objects
|
|
24
|
+
and handling exceptions during __repr__ calls (e.g. strict objects).
|
|
25
|
+
"""
|
|
26
|
+
try:
|
|
27
|
+
r = reprlib.repr(obj)
|
|
28
|
+
if len(r) > max_len:
|
|
29
|
+
return r[:max_len] + "..."
|
|
30
|
+
return r
|
|
31
|
+
except Exception:
|
|
32
|
+
return "<unrepresentable>"
|
|
33
|
+
|
|
34
|
+
def _format_args(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> str:
|
|
35
|
+
"""
|
|
36
|
+
Formats function arguments into a single string "arg1, arg2, k=v".
|
|
37
|
+
Used for the arrow label in the diagram.
|
|
38
|
+
"""
|
|
39
|
+
parts = []
|
|
40
|
+
for arg in args:
|
|
41
|
+
parts.append(_safe_repr(arg))
|
|
42
|
+
|
|
43
|
+
for k, v in kwargs.items():
|
|
44
|
+
parts.append(f"{k}={_safe_repr(v)}")
|
|
45
|
+
|
|
46
|
+
return ", ".join(parts)
|
|
47
|
+
|
|
48
|
+
def _resolve_target(func: Callable[..., Any], args: Tuple[Any, ...], target_override: Optional[str]) -> str:
|
|
49
|
+
"""
|
|
50
|
+
Determines the name of the participant (Target) for the diagram.
|
|
51
|
+
|
|
52
|
+
Resolution Priority:
|
|
53
|
+
1. **Override**: If the user explicitly provided `target="Name"`, use it.
|
|
54
|
+
2. **Instance Method**: If the first arg looks like `self` (has __class__),
|
|
55
|
+
use the class name.
|
|
56
|
+
3. **Class Method**: If the first arg is a type (cls), use the class name.
|
|
57
|
+
4. **Module Function**: Fallback to the name of the module containing the function.
|
|
58
|
+
5. **Fallback**: "Unknown".
|
|
59
|
+
"""
|
|
60
|
+
if target_override:
|
|
61
|
+
return target_override
|
|
62
|
+
|
|
63
|
+
# Heuristic: If it's a method call, args[0] is usually 'self'.
|
|
64
|
+
if args:
|
|
65
|
+
first_arg = args[0]
|
|
66
|
+
# Check if it looks like a class instance
|
|
67
|
+
# We check hasattr(__class__) to distinguish objects from primitives/containers broadly,
|
|
68
|
+
# ensuring we don't mislabel a plain list passed as first arg to a function as a "List" participant.
|
|
69
|
+
if hasattr(first_arg, "__class__") and not isinstance(first_arg, (str, int, float, bool, list, dict, type)):
|
|
70
|
+
return str(first_arg.__class__.__name__)
|
|
71
|
+
# Check if it looks like a class (cls) - e.g. @classmethod
|
|
72
|
+
if isinstance(first_arg, type):
|
|
73
|
+
return first_arg.__name__
|
|
74
|
+
|
|
75
|
+
# Fallback to module name for standalone functions
|
|
76
|
+
module = inspect.getmodule(func)
|
|
77
|
+
if module:
|
|
78
|
+
return module.__name__.split(".")[-1]
|
|
79
|
+
return "Unknown"
|
|
80
|
+
|
|
81
|
+
def _log_interaction(logger: logging.Logger,
|
|
82
|
+
source: str,
|
|
83
|
+
target: str,
|
|
84
|
+
action: str,
|
|
85
|
+
params: str,
|
|
86
|
+
trace_id: str) -> None:
|
|
87
|
+
"""
|
|
88
|
+
Logs the 'Call' event (Start of function).
|
|
89
|
+
Arrow: source -> target
|
|
90
|
+
"""
|
|
91
|
+
req_event = FlowEvent(
|
|
92
|
+
source=source, target=target,
|
|
93
|
+
action=action, message=action,
|
|
94
|
+
params=params, trace_id=trace_id
|
|
95
|
+
)
|
|
96
|
+
# The 'extra' dict is critical: the Handler will pick this up to format the Mermaid line
|
|
97
|
+
logger.info(f"{source}->{target}: {action}", extra={"flow_event": req_event})
|
|
98
|
+
|
|
99
|
+
def _log_return(logger: logging.Logger,
|
|
100
|
+
source: str,
|
|
101
|
+
target: str,
|
|
102
|
+
action: str,
|
|
103
|
+
result: Any,
|
|
104
|
+
trace_id: str) -> None:
|
|
105
|
+
"""
|
|
106
|
+
Logs the 'Return' event (End of function).
|
|
107
|
+
Arrow: target --> source (Dotted line return)
|
|
108
|
+
|
|
109
|
+
Note: 'source' here is the original caller, 'target' is the callee.
|
|
110
|
+
So the return arrow goes from target back to source.
|
|
111
|
+
"""
|
|
112
|
+
result_str = _safe_repr(result)
|
|
113
|
+
resp_event = FlowEvent(
|
|
114
|
+
source=target, target=source,
|
|
115
|
+
action=action, message="Return",
|
|
116
|
+
is_return=True, result=result_str, trace_id=trace_id
|
|
117
|
+
)
|
|
118
|
+
logger.info(f"{target}->{source}: Return", extra={"flow_event": resp_event})
|
|
119
|
+
|
|
120
|
+
def _log_error(logger: logging.Logger,
|
|
121
|
+
source: str,
|
|
122
|
+
target: str,
|
|
123
|
+
action: str,
|
|
124
|
+
error: Exception,
|
|
125
|
+
trace_id: str) -> None:
|
|
126
|
+
"""
|
|
127
|
+
Logs an 'Error' event if the function raises an exception.
|
|
128
|
+
Arrow: target -x source (Error return)
|
|
129
|
+
"""
|
|
130
|
+
err_event = FlowEvent(
|
|
131
|
+
source=target, target=source, action=action,
|
|
132
|
+
message=str(error), is_return=True, is_error=True, error_message=str(error),
|
|
133
|
+
trace_id=trace_id
|
|
134
|
+
)
|
|
135
|
+
logger.error(f"{target}-x{source}: Error", extra={"flow_event": err_event})
|
|
136
|
+
|
|
137
|
+
@overload
|
|
138
|
+
def trace_interaction(func: F) -> F:
|
|
139
|
+
...
|
|
140
|
+
|
|
141
|
+
@overload
|
|
142
|
+
def trace_interaction(
|
|
143
|
+
*,
|
|
144
|
+
source: Optional[str] = None,
|
|
145
|
+
target: Optional[str] = None,
|
|
146
|
+
action: Optional[str] = None
|
|
147
|
+
) -> Callable[[F], F]:
|
|
148
|
+
...
|
|
149
|
+
|
|
150
|
+
def trace_interaction(
|
|
151
|
+
func: Optional[F] = None,
|
|
152
|
+
*,
|
|
153
|
+
source: Optional[str] = None,
|
|
154
|
+
target: Optional[str] = None,
|
|
155
|
+
action: Optional[str] = None
|
|
156
|
+
) -> Union[F, Callable[[F], F]]:
|
|
157
|
+
"""
|
|
158
|
+
Main Decorator for tracing function execution in Mermaid diagrams.
|
|
159
|
+
|
|
160
|
+
It supports two modes of operation:
|
|
161
|
+
1. **Simple**: `@trace` (No arguments)
|
|
162
|
+
2. **Configured**: `@trace(action="Login", target="AuthService")`
|
|
163
|
+
|
|
164
|
+
Args:
|
|
165
|
+
func: The function being decorated (automatically passed in simple mode).
|
|
166
|
+
source: Explicit name of the caller participant (rarely used, usually inferred from Context).
|
|
167
|
+
target: Explicit name of the callee participant (overrides automatic resolution).
|
|
168
|
+
action: Label for the arrow (defaults to function name).
|
|
169
|
+
"""
|
|
170
|
+
|
|
171
|
+
# Mode 1: @trace used without parentheses
|
|
172
|
+
if func is not None and callable(func):
|
|
173
|
+
return _create_decorator(func, source, target, action)
|
|
174
|
+
|
|
175
|
+
# Mode 2: @trace(...) used with arguments -> returns a factory
|
|
176
|
+
def factory(f: F) -> F:
|
|
177
|
+
return _create_decorator(f, source, target, action)
|
|
178
|
+
return factory
|
|
179
|
+
|
|
180
|
+
def _create_decorator(
|
|
181
|
+
func: F,
|
|
182
|
+
source: Optional[str],
|
|
183
|
+
target: Optional[str],
|
|
184
|
+
action: Optional[str]
|
|
185
|
+
) -> F:
|
|
186
|
+
"""
|
|
187
|
+
Constructs the actual wrapper function.
|
|
188
|
+
Handles both synchronous and asynchronous functions.
|
|
189
|
+
"""
|
|
190
|
+
|
|
191
|
+
# Pre-calculate static metadata to save time at runtime
|
|
192
|
+
if action is None:
|
|
193
|
+
action = func.__name__.replace("_", " ").title()
|
|
194
|
+
|
|
195
|
+
@functools.wraps(func)
|
|
196
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
197
|
+
"""Sync function wrapper."""
|
|
198
|
+
# 1. Resolve Context
|
|
199
|
+
# 'source' is who called us (from Context). 'target' is who we are (resolved from self/cls).
|
|
200
|
+
current_source = source or LogContext.current_participant()
|
|
201
|
+
trace_id = LogContext.current_trace_id()
|
|
202
|
+
current_target = _resolve_target(func, args, target)
|
|
203
|
+
|
|
204
|
+
logger = get_flow_logger()
|
|
205
|
+
params_str = _format_args(args, kwargs)
|
|
206
|
+
|
|
207
|
+
# 2. Log Request (Start of block)
|
|
208
|
+
_log_interaction(logger, current_source, current_target, action, params_str, trace_id)
|
|
209
|
+
|
|
210
|
+
# 3. Execute with New Context
|
|
211
|
+
# We push 'current_target' as the NEW 'participant' (source) for any internal calls.
|
|
212
|
+
with LogContext.scope({"participant": current_target, "trace_id": trace_id}):
|
|
213
|
+
try:
|
|
214
|
+
result = func(*args, **kwargs)
|
|
215
|
+
# 4. Log Success Return
|
|
216
|
+
_log_return(logger, current_source, current_target, action, result, trace_id)
|
|
217
|
+
return result
|
|
218
|
+
except Exception as e:
|
|
219
|
+
# 5. Log Error Return
|
|
220
|
+
_log_error(logger, current_source, current_target, action, e, trace_id)
|
|
221
|
+
raise
|
|
222
|
+
|
|
223
|
+
@functools.wraps(func)
|
|
224
|
+
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
225
|
+
"""Async function wrapper (coroutine)."""
|
|
226
|
+
current_source = source or LogContext.current_participant()
|
|
227
|
+
trace_id = LogContext.current_trace_id()
|
|
228
|
+
current_target = _resolve_target(func, args, target)
|
|
229
|
+
|
|
230
|
+
logger = get_flow_logger()
|
|
231
|
+
params_str = _format_args(args, kwargs)
|
|
232
|
+
|
|
233
|
+
# 2. Log Request (Start of block)
|
|
234
|
+
_log_interaction(logger, current_source, current_target, action, params_str, trace_id)
|
|
235
|
+
|
|
236
|
+
# Use async context manager (ascope) to ensure context propagates correctly across awaits
|
|
237
|
+
async with LogContext.ascope({"participant": current_target, "trace_id": trace_id}):
|
|
238
|
+
try:
|
|
239
|
+
result = await func(*args, **kwargs)
|
|
240
|
+
_log_return(logger, current_source, current_target, action, result, trace_id)
|
|
241
|
+
return result
|
|
242
|
+
except Exception as e:
|
|
243
|
+
_log_error(logger, current_source, current_target, action, e, trace_id)
|
|
244
|
+
raise
|
|
245
|
+
|
|
246
|
+
# Detect if the wrapped function is a coroutine to choose the right wrapper
|
|
247
|
+
if inspect.iscoroutinefunction(func):
|
|
248
|
+
return cast(F, async_wrapper)
|
|
249
|
+
return cast(F, wrapper)
|
|
250
|
+
|
|
251
|
+
# Alias for easy import
|
|
252
|
+
trace = trace_interaction
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
from dataclasses import dataclass, field
|
|
2
|
+
import time
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
@dataclass
|
|
6
|
+
class FlowEvent:
|
|
7
|
+
"""
|
|
8
|
+
Represents a single interaction or step in the execution flow.
|
|
9
|
+
|
|
10
|
+
This data structure acts as the intermediate representation (IR) between
|
|
11
|
+
runtime code execution and the final Mermaid diagram output. Each instance
|
|
12
|
+
corresponds directly to one arrow or note in the sequence diagram.
|
|
13
|
+
|
|
14
|
+
The fields map to Mermaid syntax components as follows:
|
|
15
|
+
`source` -> `target`: `message`
|
|
16
|
+
|
|
17
|
+
Attributes:
|
|
18
|
+
source (str):
|
|
19
|
+
The name of the participant initiating the action (the "Caller").
|
|
20
|
+
In Mermaid: The participant on the LEFT side of the arrow.
|
|
21
|
+
|
|
22
|
+
target (str):
|
|
23
|
+
The name of the participant receiving the action (the "Callee").
|
|
24
|
+
In Mermaid: The participant on the RIGHT side of the arrow.
|
|
25
|
+
|
|
26
|
+
action (str):
|
|
27
|
+
A short, human-readable name for the operation (e.g., function name).
|
|
28
|
+
Used for grouping or filtering logs, but often redundant with message.
|
|
29
|
+
|
|
30
|
+
message (str):
|
|
31
|
+
The actual text label displayed on the diagram arrow.
|
|
32
|
+
Example: "getUser(id=1)" or "Return: User(name='Alice')".
|
|
33
|
+
|
|
34
|
+
timestamp (float):
|
|
35
|
+
Unix timestamp (seconds) of when the event occurred.
|
|
36
|
+
Used for ordering events if logs are processed asynchronously,
|
|
37
|
+
though Mermaid sequence diagrams primarily rely on line order.
|
|
38
|
+
|
|
39
|
+
trace_id (str):
|
|
40
|
+
Unique identifier for the trace session.
|
|
41
|
+
Allows filtering multiple concurrent traces from a single log file
|
|
42
|
+
to generate separate diagrams for separate requests.
|
|
43
|
+
|
|
44
|
+
is_return (bool):
|
|
45
|
+
Flag indicating if this is a response arrow.
|
|
46
|
+
If True, the arrow is drawn as a dotted line (`-->`) in Mermaid.
|
|
47
|
+
If False, it is a solid line (`->`) representing a call.
|
|
48
|
+
|
|
49
|
+
is_error (bool):
|
|
50
|
+
Flag indicating if an exception occurred.
|
|
51
|
+
If True, the arrow might be styled differently (e.g., `-x`) to show failure.
|
|
52
|
+
|
|
53
|
+
error_message (Optional[str]):
|
|
54
|
+
Detailed error text if `is_error` is True.
|
|
55
|
+
Can be added as a note or included in the arrow label.
|
|
56
|
+
|
|
57
|
+
params (Optional[str]):
|
|
58
|
+
Stringified representation of function arguments.
|
|
59
|
+
Captured only for request events (call start).
|
|
60
|
+
|
|
61
|
+
result (Optional[str]):
|
|
62
|
+
Stringified representation of the return value.
|
|
63
|
+
Captured only for return events (call end).
|
|
64
|
+
"""
|
|
65
|
+
source: str
|
|
66
|
+
target: str
|
|
67
|
+
action: str
|
|
68
|
+
message: str
|
|
69
|
+
trace_id: str
|
|
70
|
+
timestamp: float = field(default_factory=time.time)
|
|
71
|
+
is_return: bool = False
|
|
72
|
+
is_error: bool = False
|
|
73
|
+
error_message: Optional[str] = None
|
|
74
|
+
params: Optional[str] = None
|
|
75
|
+
result: Optional[str] = None
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import re
|
|
3
|
+
from typing import Optional
|
|
4
|
+
from .events import FlowEvent
|
|
5
|
+
|
|
6
|
+
class MermaidFormatter(logging.Formatter):
|
|
7
|
+
"""
|
|
8
|
+
Custom formatter to convert FlowEvents into Mermaid sequence diagram syntax.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
def format(self, record: logging.LogRecord) -> str:
|
|
12
|
+
# 1. Retrieve the FlowEvent
|
|
13
|
+
event: Optional[FlowEvent] = getattr(record, 'flow_event', None)
|
|
14
|
+
|
|
15
|
+
if not event:
|
|
16
|
+
# Fallback for standard logs if they accidentally reach this handler
|
|
17
|
+
return super().format(record)
|
|
18
|
+
|
|
19
|
+
# 2. Convert event to Mermaid line
|
|
20
|
+
return self._to_mermaid_line(event)
|
|
21
|
+
|
|
22
|
+
def _to_mermaid_line(self, event: FlowEvent) -> str:
|
|
23
|
+
"""
|
|
24
|
+
Converts a FlowEvent into a Mermaid syntax string.
|
|
25
|
+
"""
|
|
26
|
+
# Sanitize participant names
|
|
27
|
+
src = self._sanitize(event.source)
|
|
28
|
+
tgt = self._sanitize(event.target)
|
|
29
|
+
|
|
30
|
+
# Determine arrow type
|
|
31
|
+
# ->> : Solid line with arrowhead (synchronous call)
|
|
32
|
+
# -->> : Dotted line with arrowhead (return)
|
|
33
|
+
# --x : Dotted line with cross (error)
|
|
34
|
+
arrow = "-->>" if event.is_return else "->>"
|
|
35
|
+
|
|
36
|
+
if event.is_error:
|
|
37
|
+
arrow = "--x"
|
|
38
|
+
msg = f"Error: {event.error_message}"
|
|
39
|
+
elif event.is_return:
|
|
40
|
+
msg = f"Return: {event.result}" if event.result else "Return"
|
|
41
|
+
else:
|
|
42
|
+
msg = f"{event.message}({event.params})" if event.params else event.message
|
|
43
|
+
|
|
44
|
+
# Optional: Add note or group if trace_id changes (not implemented in single line format)
|
|
45
|
+
# For now, we just output the interaction.
|
|
46
|
+
|
|
47
|
+
# Escape message for Mermaid safety
|
|
48
|
+
msg = self._escape_message(msg)
|
|
49
|
+
|
|
50
|
+
return f"{src}{arrow}{tgt}: {msg}"
|
|
51
|
+
|
|
52
|
+
def _sanitize(self, name: str) -> str:
|
|
53
|
+
"""
|
|
54
|
+
Sanitizes participant names to be valid Mermaid identifiers.
|
|
55
|
+
Allows alphanumeric and underscores. Replaces others.
|
|
56
|
+
"""
|
|
57
|
+
# Replace any non-alphanumeric character (except underscore) with underscore
|
|
58
|
+
clean_name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
|
|
59
|
+
# Ensure it doesn't start with a digit (Mermaid doesn't like that sometimes, though often okay)
|
|
60
|
+
if clean_name and clean_name[0].isdigit():
|
|
61
|
+
clean_name = "_" + clean_name
|
|
62
|
+
return clean_name
|
|
63
|
+
|
|
64
|
+
def _escape_message(self, msg: str) -> str:
|
|
65
|
+
"""
|
|
66
|
+
Escapes special characters in the message text.
|
|
67
|
+
Mermaid messages can contain most chars, but : and newlines can be tricky.
|
|
68
|
+
"""
|
|
69
|
+
# Replace newlines with <br/> for Mermaid display
|
|
70
|
+
msg = msg.replace('\n', '<br/>')
|
|
71
|
+
# We might want to escape other chars if needed, but usually text after : is forgiving.
|
|
72
|
+
return msg
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
|
|
4
|
+
class MermaidFileHandler(logging.FileHandler):
|
|
5
|
+
"""
|
|
6
|
+
A custom logging handler that writes `FlowEvent` objects to a Mermaid (.mmd) file.
|
|
7
|
+
|
|
8
|
+
Strategy & Optimization:
|
|
9
|
+
1. **Inheritance**: Inherits from `logging.FileHandler` to leverage robust,
|
|
10
|
+
thread-safe file writing capabilities (locking, buffering) provided by the stdlib.
|
|
11
|
+
2. **Header Management**: Automatically handles the Mermaid file header
|
|
12
|
+
(`sequenceDiagram`, `title`, `autonumber`) to ensure the output file
|
|
13
|
+
is a valid Mermaid document. It smartly detects if the file is new or
|
|
14
|
+
being appended to.
|
|
15
|
+
3. **Deferred Formatting**: The actual string conversion happens in the `emit`
|
|
16
|
+
method (via the formatter), keeping the handler focused on I/O.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
def __init__(self, filename: str, title: str = "Log Flow", mode: str = 'a', encoding: str = 'utf-8', delay: bool = False):
|
|
20
|
+
"""
|
|
21
|
+
Initialize the handler.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
filename (str): The path to the output .mmd file.
|
|
25
|
+
title (str): The title of the Mermaid diagram (written in the header).
|
|
26
|
+
mode (str): File open mode. 'w' (overwrite) or 'a' (append).
|
|
27
|
+
encoding (str): File encoding. Defaults to 'utf-8'.
|
|
28
|
+
delay (bool): If True, file opening is deferred until the first call to emit.
|
|
29
|
+
Useful to avoid creating empty files if no logs occur.
|
|
30
|
+
"""
|
|
31
|
+
# Ensure directory exists to prevent FileNotFoundError on open
|
|
32
|
+
os.makedirs(os.path.dirname(os.path.abspath(filename)) or ".", exist_ok=True)
|
|
33
|
+
|
|
34
|
+
# Header Strategy:
|
|
35
|
+
# We need to write the "sequenceDiagram" preamble ONLY if:
|
|
36
|
+
# 1. We are overwriting the file (mode='w').
|
|
37
|
+
# 2. We are appending (mode='a'), but the file doesn't exist or is empty.
|
|
38
|
+
should_write_header = False
|
|
39
|
+
if mode == 'w':
|
|
40
|
+
should_write_header = True
|
|
41
|
+
elif mode == 'a':
|
|
42
|
+
if not os.path.exists(filename) or os.path.getsize(filename) == 0:
|
|
43
|
+
should_write_header = True
|
|
44
|
+
|
|
45
|
+
# Initialize standard FileHandler (opens the file unless delay=True)
|
|
46
|
+
super().__init__(filename, mode, encoding, delay)
|
|
47
|
+
self.title = title
|
|
48
|
+
|
|
49
|
+
# Write header immediately if needed.
|
|
50
|
+
if should_write_header:
|
|
51
|
+
self._write_header()
|
|
52
|
+
|
|
53
|
+
def _write_header(self) -> None:
|
|
54
|
+
"""
|
|
55
|
+
Writes the initial Mermaid syntax lines.
|
|
56
|
+
|
|
57
|
+
This setup is required for Mermaid JS or Live Editor to render the diagram.
|
|
58
|
+
"""
|
|
59
|
+
# We use the stream directly if available, or open momentarily if delayed
|
|
60
|
+
if self.stream:
|
|
61
|
+
self.stream.write("sequenceDiagram\n")
|
|
62
|
+
self.stream.write(f" title {self.title}\n")
|
|
63
|
+
self.stream.write(" autonumber\n")
|
|
64
|
+
self.flush()
|
|
65
|
+
else:
|
|
66
|
+
# Handling 'delay=True' case:
|
|
67
|
+
# If the file isn't open yet, we temporarily open it just to write the header.
|
|
68
|
+
# This ensures the file is valid even if the application crashes before the first log.
|
|
69
|
+
with open(self.baseFilename, self.mode, encoding=self.encoding) as f:
|
|
70
|
+
f.write("sequenceDiagram\n")
|
|
71
|
+
f.write(f" title {self.title}\n")
|
|
72
|
+
f.write(" autonumber\n")
|
|
73
|
+
|
|
74
|
+
def emit(self, record: logging.LogRecord) -> None:
|
|
75
|
+
"""
|
|
76
|
+
Process a log record.
|
|
77
|
+
|
|
78
|
+
Optimization:
|
|
79
|
+
- Checks for `flow_event` attribute first. This allows this handler
|
|
80
|
+
to be attached to the root logger without processing irrelevant system logs.
|
|
81
|
+
- Delegates the actual writing to `super().emit()`, which handles
|
|
82
|
+
thread locking and stream flushing.
|
|
83
|
+
"""
|
|
84
|
+
# Only process records that contain our structured FlowEvent data
|
|
85
|
+
if hasattr(record, 'flow_event'):
|
|
86
|
+
super().emit(record)
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
from typing import Any, TYPE_CHECKING
|
|
2
|
+
import time
|
|
3
|
+
import uuid
|
|
4
|
+
|
|
5
|
+
from ..core.events import FlowEvent
|
|
6
|
+
from ..core.context import LogContext
|
|
7
|
+
from ..core.decorators import get_flow_logger
|
|
8
|
+
|
|
9
|
+
if TYPE_CHECKING:
|
|
10
|
+
from fastapi import Request, Response
|
|
11
|
+
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
|
12
|
+
else:
|
|
13
|
+
try:
|
|
14
|
+
from fastapi import Request, Response
|
|
15
|
+
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
|
16
|
+
except ImportError:
|
|
17
|
+
# Handle the case where FastAPI/Starlette are not installed.
|
|
18
|
+
# We define dummy types to prevent NameErrors at import time,
|
|
19
|
+
# but instantiation will fail explicitly in __init__.
|
|
20
|
+
BaseHTTPMiddleware = object # type: ignore[misc,assignment]
|
|
21
|
+
Request = Any # type: ignore[assignment]
|
|
22
|
+
Response = Any # type: ignore[assignment]
|
|
23
|
+
RequestResponseEndpoint = Any # type: ignore[assignment]
|
|
24
|
+
|
|
25
|
+
class MermaidTraceMiddleware(BaseHTTPMiddleware):
|
|
26
|
+
"""
|
|
27
|
+
FastAPI Middleware to trace HTTP requests as interactions in the sequence diagram.
|
|
28
|
+
|
|
29
|
+
This middleware acts as the entry point for tracing a web request. It:
|
|
30
|
+
1. Identifies the client (Source).
|
|
31
|
+
2. Logs the incoming request.
|
|
32
|
+
3. Initializes the `LogContext` for the request lifecycle.
|
|
33
|
+
4. Logs the response or error.
|
|
34
|
+
"""
|
|
35
|
+
def __init__(self, app: Any, app_name: str = "FastAPI"):
|
|
36
|
+
"""
|
|
37
|
+
Initialize the middleware.
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
app: The FastAPI application instance.
|
|
41
|
+
app_name: The name of this service to appear in the diagram (e.g., "UserAPI").
|
|
42
|
+
"""
|
|
43
|
+
if BaseHTTPMiddleware is object: # type: ignore[comparison-overlap]
|
|
44
|
+
raise ImportError("FastAPI/Starlette is required to use MermaidTraceMiddleware")
|
|
45
|
+
super().__init__(app)
|
|
46
|
+
self.app_name = app_name
|
|
47
|
+
|
|
48
|
+
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
|
49
|
+
"""
|
|
50
|
+
Intercepts the incoming request.
|
|
51
|
+
|
|
52
|
+
Args:
|
|
53
|
+
request (Request): The incoming HTTP request.
|
|
54
|
+
call_next (Callable): The function to call the next middleware or endpoint.
|
|
55
|
+
|
|
56
|
+
Returns:
|
|
57
|
+
Response: The HTTP response.
|
|
58
|
+
"""
|
|
59
|
+
# 1. Determine Source (Client)
|
|
60
|
+
# Try to get a specific ID from headers (useful for distributed tracing),
|
|
61
|
+
# otherwise fallback to "Client".
|
|
62
|
+
source = request.headers.get("X-Source", "Client")
|
|
63
|
+
|
|
64
|
+
# Determine Trace ID
|
|
65
|
+
# Check X-Trace-ID header or generate new UUID
|
|
66
|
+
trace_id = request.headers.get("X-Trace-ID") or str(uuid.uuid4())
|
|
67
|
+
|
|
68
|
+
# 2. Determine Action
|
|
69
|
+
# Format: "METHOD /path" (e.g., "GET /users")
|
|
70
|
+
action = f"{request.method} {request.url.path}"
|
|
71
|
+
|
|
72
|
+
logger = get_flow_logger()
|
|
73
|
+
|
|
74
|
+
# 3. Log Request (Source -> App)
|
|
75
|
+
req_event = FlowEvent(
|
|
76
|
+
source=source,
|
|
77
|
+
target=self.app_name,
|
|
78
|
+
action=action,
|
|
79
|
+
message=action,
|
|
80
|
+
params=f"query={request.query_params}" if request.query_params else None,
|
|
81
|
+
trace_id=trace_id
|
|
82
|
+
)
|
|
83
|
+
logger.info(f"{source}->{self.app_name}: {action}", extra={"flow_event": req_event})
|
|
84
|
+
|
|
85
|
+
# 4. Set Context and Process Request
|
|
86
|
+
# We set the current participant to the app name.
|
|
87
|
+
# `ascope` ensures this context applies to all code running within `call_next`.
|
|
88
|
+
async with LogContext.ascope({"participant": self.app_name, "trace_id": trace_id}):
|
|
89
|
+
start_time = time.time()
|
|
90
|
+
try:
|
|
91
|
+
# Pass control to the application
|
|
92
|
+
response = await call_next(request)
|
|
93
|
+
|
|
94
|
+
# 5. Log Response (App -> Source)
|
|
95
|
+
duration = (time.time() - start_time) * 1000
|
|
96
|
+
resp_event = FlowEvent(
|
|
97
|
+
source=self.app_name,
|
|
98
|
+
target=source,
|
|
99
|
+
action=action,
|
|
100
|
+
message="Return",
|
|
101
|
+
is_return=True,
|
|
102
|
+
result=f"{response.status_code} ({duration:.1f}ms)",
|
|
103
|
+
trace_id=trace_id
|
|
104
|
+
)
|
|
105
|
+
logger.info(f"{self.app_name}->{source}: Return", extra={"flow_event": resp_event})
|
|
106
|
+
return response
|
|
107
|
+
|
|
108
|
+
except Exception as e:
|
|
109
|
+
# 6. Log Error (App --x Source)
|
|
110
|
+
# This captures unhandled exceptions that bubble up to the middleware
|
|
111
|
+
err_event = FlowEvent(
|
|
112
|
+
source=self.app_name,
|
|
113
|
+
target=source,
|
|
114
|
+
action=action,
|
|
115
|
+
message=str(e),
|
|
116
|
+
is_return=True,
|
|
117
|
+
is_error=True,
|
|
118
|
+
error_message=str(e),
|
|
119
|
+
trace_id=trace_id
|
|
120
|
+
)
|
|
121
|
+
logger.error(f"{self.app_name}-x{source}: Error", extra={"flow_event": err_event})
|
|
122
|
+
raise
|
mermaid_trace/py.typed
ADDED
|
File without changes
|
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mermaid-trace
|
|
3
|
+
Version: 0.3.1
|
|
4
|
+
Summary: Visualize your Python code execution flow as Mermaid Sequence Diagrams.
|
|
5
|
+
Project-URL: Documentation, https://github.com/xt765/mermaid-trace#readme
|
|
6
|
+
Project-URL: Changelog, https://github.com/xt765/mermaid-trace/blob/main/docs/en/CHANGELOG.md
|
|
7
|
+
Project-URL: Issues, https://github.com/xt765/mermaid-trace/issues
|
|
8
|
+
Project-URL: Source, https://github.com/xt765/mermaid-trace
|
|
9
|
+
Author-email: xt765 <xt765@foxmail.com>
|
|
10
|
+
License: MIT License
|
|
11
|
+
|
|
12
|
+
Copyright (c) 2026 xt765
|
|
13
|
+
|
|
14
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
15
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
16
|
+
in the Software without restriction, including without limitation the rights
|
|
17
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
18
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
19
|
+
furnished to do so, subject to the following conditions:
|
|
20
|
+
|
|
21
|
+
The above copyright notice and this permission notice shall be included in all
|
|
22
|
+
copies or substantial portions of the Software.
|
|
23
|
+
|
|
24
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
25
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
26
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
27
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
28
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
29
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
30
|
+
SOFTWARE.
|
|
31
|
+
License-File: LICENSE
|
|
32
|
+
Keywords: asyncio,logging,mermaid,mermaid-trace,sequence-diagram,trace,visualization
|
|
33
|
+
Classifier: Development Status :: 4 - Beta
|
|
34
|
+
Classifier: Intended Audience :: Developers
|
|
35
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
36
|
+
Classifier: Operating System :: OS Independent
|
|
37
|
+
Classifier: Programming Language :: Python
|
|
38
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
39
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
40
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
41
|
+
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
42
|
+
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
43
|
+
Classifier: Topic :: Software Development :: Debuggers
|
|
44
|
+
Classifier: Topic :: System :: Logging
|
|
45
|
+
Requires-Python: >=3.10
|
|
46
|
+
Requires-Dist: typing-extensions>=4.0.0
|
|
47
|
+
Requires-Dist: watchdog>=2.0.0
|
|
48
|
+
Provides-Extra: all
|
|
49
|
+
Requires-Dist: fastapi>=0.100.0; extra == 'all'
|
|
50
|
+
Provides-Extra: dev
|
|
51
|
+
Requires-Dist: fastapi>=0.100.0; extra == 'dev'
|
|
52
|
+
Requires-Dist: httpx; extra == 'dev'
|
|
53
|
+
Requires-Dist: mypy; extra == 'dev'
|
|
54
|
+
Requires-Dist: pytest; extra == 'dev'
|
|
55
|
+
Requires-Dist: pytest-asyncio; extra == 'dev'
|
|
56
|
+
Requires-Dist: pytest-cov; extra == 'dev'
|
|
57
|
+
Requires-Dist: ruff; extra == 'dev'
|
|
58
|
+
Provides-Extra: fastapi
|
|
59
|
+
Requires-Dist: fastapi>=0.100.0; extra == 'fastapi'
|
|
60
|
+
Description-Content-Type: text/markdown
|
|
61
|
+
|
|
62
|
+
# MermaidTrace: The Python Logger That Draws Diagrams
|
|
63
|
+
|
|
64
|
+
[](https://pypi.org/project/mermaid-trace/)
|
|
65
|
+
[](https://pypi.org/project/mermaid-trace/)
|
|
66
|
+
[](LICENSE)
|
|
67
|
+
[](https://github.com/xt765/mermaid-trace/actions/workflows/ci.yml)
|
|
68
|
+
[](https://codecov.io/gh/xt765/mermaid-trace)
|
|
69
|
+
|
|
70
|
+
**Stop reading logs. Start watching them.**
|
|
71
|
+
|
|
72
|
+
MermaidTrace is a specialized logging tool that automatically generates [Mermaid JS](https://mermaid.js.org/) sequence diagrams from your code execution. It's perfect for visualizing complex business logic, microservice interactions, or asynchronous flows.
|
|
73
|
+
|
|
74
|
+
## β¨ Features
|
|
75
|
+
|
|
76
|
+
- **Decorator-Driven**: Just add `@trace` or `@trace_interaction` to your functions.
|
|
77
|
+
- **Auto-Diagramming**: Generates `.mmd` files that can be viewed in VS Code, GitHub, or Mermaid Live Editor.
|
|
78
|
+
- **Async Support**: Works seamlessly with `asyncio` coroutines.
|
|
79
|
+
- **Context Inference**: Automatically tracks nested calls and infers `source` participants using `contextvars`.
|
|
80
|
+
- **FastAPI Integration**: Includes middleware for zero-config HTTP request tracing.
|
|
81
|
+
- **CLI Tool**: Built-in viewer to preview diagrams in your browser.
|
|
82
|
+
|
|
83
|
+
## π Quick Start
|
|
84
|
+
|
|
85
|
+
### Installation
|
|
86
|
+
|
|
87
|
+
```bash
|
|
88
|
+
pip install mermaid-trace
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
### Basic Usage
|
|
92
|
+
|
|
93
|
+
```python
|
|
94
|
+
from mermaid_trace import trace, configure_flow
|
|
95
|
+
import time
|
|
96
|
+
|
|
97
|
+
# 1. Configure output
|
|
98
|
+
configure_flow("my_flow.mmd")
|
|
99
|
+
|
|
100
|
+
# 2. Add decorators
|
|
101
|
+
@trace(source="Client", target="PaymentService", action="Process Payment")
|
|
102
|
+
def process_payment(amount):
|
|
103
|
+
if check_balance(amount):
|
|
104
|
+
return "Success"
|
|
105
|
+
return "Failed"
|
|
106
|
+
|
|
107
|
+
@trace(source="PaymentService", target="Database", action="Check Balance")
|
|
108
|
+
def check_balance(amount):
|
|
109
|
+
return True
|
|
110
|
+
|
|
111
|
+
# 3. Run your code
|
|
112
|
+
process_payment(100)
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
### Nested Calls (Context Inference)
|
|
116
|
+
|
|
117
|
+
You don't need to specify `source` every time. MermaidTrace infers it from the current context.
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
@trace(source="Client", target="API")
|
|
121
|
+
def main():
|
|
122
|
+
# Inside here, current participant is "API"
|
|
123
|
+
service_call()
|
|
124
|
+
|
|
125
|
+
@trace(target="Service") # source inferred as "API"
|
|
126
|
+
def service_call():
|
|
127
|
+
pass
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
### FastAPI Integration
|
|
131
|
+
|
|
132
|
+
```python
|
|
133
|
+
from fastapi import FastAPI
|
|
134
|
+
from mermaid_trace.integrations.fastapi import MermaidTraceMiddleware
|
|
135
|
+
|
|
136
|
+
app = FastAPI()
|
|
137
|
+
app.add_middleware(MermaidTraceMiddleware, app_name="MyAPI")
|
|
138
|
+
|
|
139
|
+
@app.get("/")
|
|
140
|
+
async def root():
|
|
141
|
+
return {"message": "Hello World"}
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
### CLI Viewer
|
|
145
|
+
|
|
146
|
+
Visualize your generated `.mmd` files instantly:
|
|
147
|
+
|
|
148
|
+
```bash
|
|
149
|
+
mermaid-trace serve my_flow.mmd
|
|
150
|
+
```
|
|
151
|
+
|
|
152
|
+
## π Documentation
|
|
153
|
+
|
|
154
|
+
- [English Documentation](docs/en/README.md)
|
|
155
|
+
- [δΈζζζ‘£](README_CN.md)
|
|
156
|
+
|
|
157
|
+
## π€ Contributing
|
|
158
|
+
|
|
159
|
+
We welcome contributions! Please see [CONTRIBUTING.md](docs/en/CONTRIBUTING.md) for details.
|
|
160
|
+
|
|
161
|
+
## π License
|
|
162
|
+
|
|
163
|
+
MIT
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
mermaid_trace/__init__.py,sha256=yEtdgB0Cl2qAuhGnU0zp9_YANzT3ZIrDsnQu-i3qUAc,2741
|
|
2
|
+
mermaid_trace/cli.py,sha256=_6Bm6oGIUQFbb-tr4yDoZIa9_5PQ-qjftCd7JNUm4F8,7033
|
|
3
|
+
mermaid_trace/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
4
|
+
mermaid_trace/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
5
|
+
mermaid_trace/core/context.py,sha256=H2aXivdarZgpvfAdATRWAAPpf7KHSNqhoJOepd1FK_E,6711
|
|
6
|
+
mermaid_trace/core/decorators.py,sha256=vjTnqxOCbIZXBahpgvXHBp3BFIjnIRk15O2jR9V8Ls4,9274
|
|
7
|
+
mermaid_trace/core/events.py,sha256=SG-S0hZYNa_gEpwB4Cuaahwj9WAMdzS7xXIpThoZzGQ,2999
|
|
8
|
+
mermaid_trace/core/formatter.py,sha256=WKXQeNj4l-LvbF2_Jz8FS3dyWmvtD3CHX2bQKj8p7D0,2765
|
|
9
|
+
mermaid_trace/handlers/mermaid_handler.py,sha256=dZ_T57qWUQlVNWWH5se52mB4ZioECOFwGJF52V-OOps,3931
|
|
10
|
+
mermaid_trace/integrations/fastapi.py,sha256=DeBodoPxXz2nLUIAiqmdVfFrwNVo2P0jhe7rcdF60u4,4991
|
|
11
|
+
mermaid_trace-0.3.1.dist-info/METADATA,sha256=rRBCzhTbof4unI0MPJkj-eHzKTazsodNpdfHpzDPDi0,6181
|
|
12
|
+
mermaid_trace-0.3.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
|
|
13
|
+
mermaid_trace-0.3.1.dist-info/entry_points.txt,sha256=WS57KT_870v0A4B87QDjQUqJcddMQxbCQyYeczDAX34,57
|
|
14
|
+
mermaid_trace-0.3.1.dist-info/licenses/LICENSE,sha256=BrBog1Etiq9PdWy0SVQNVByIMD9ss4Edz-R0oXt49zA,1062
|
|
15
|
+
mermaid_trace-0.3.1.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 xt765
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|