msb-public
/
diagrams
mirror of https://github.com/mingrammer/diagrams
1
0
Code Issues Projects Releases Wiki Activity

chore: fix pre-commit hook failures and formatting

Browse Source
pull/1203/head
Carlos Barbero 7 months ago
parent 57433dc2e2
commit adb8661cb7
5 changed files with 64 additions and 51 deletions
Show Stats Download Patch File Download Diff File

2
.isort.cfg
Unescape View File

@ -2,4 +2,4 @@
line_length = 120
multi_line_output = 3
include_trailing_comma = True
known_third_party = graphviz,jinja2
known_third_party = graphviz,inspection,jinja2,mcp

2
mcp-server/GEMINI.md
Unescape View File

@ -78,7 +78,7 @@ The server will define three main tools. It should use the `mcp.server.fastmcp`
- `code` (string, required): The Python code using `diagrams` DSL.
- `filename` (string, optional): Desired output filename.
* **Returns**:
* **Returns**:
- `status`: "success" or "error"
- `message`: Path to file or error message.

10
mcp-server/README.md
Unescape View File

@ -14,12 +14,12 @@ graph TD
subgraph Container ["Docker Container"]
MCPServer["MCP Server (server.py)"]
NodeRegistry["(Node Registry)"]
subgraph Logic ["Core Logic"]
Inspector["src/inspection.py"]
Executor["generate_diagram"]
end
subgraph Libs ["Dependencies"]
DiagramsLib["diagrams package"]
Graphviz["Graphviz Binary"]
@ -34,12 +34,12 @@ graph TD
%% Tool Flows
Client -- "list_icons()" --> MCPServer
MCPServer -- "Query" --> NodeRegistry
Client -- "generate_diagram(code)" --> MCPServer
MCPServer -- "Pass Code" --> Executor
Executor -- "exec()" --> DiagramsLib
DiagramsLib -- "Render" --> Graphviz
%% Output
Graphviz -- "Generates PNG" --> Executor
Executor -- "Writes File (Volume Mount)" --> HostFS
@ -109,7 +109,7 @@ Generates a diagram from Python code.
```python
from diagrams import Diagram
from diagrams.aws.compute import EC2
with Diagram("Simple", show=False):
EC2("web")
```

25
mcp-server/src/inspection.py
Unescape View File

@ -1,15 +1,17 @@
import pkgutil
import importlib
import inspect
import pkgutil
import sys
from collections import defaultdict
import diagrams
from diagrams import Node
def get_all_nodes():
"""
Dynamically inspects the diagrams package and returns a dictionary of all available Nodes.
Returns:
dict: A nested dictionary structure:
{
@ -20,7 +22,7 @@ def get_all_nodes():
"""
# Initialize the structure
icons = defaultdict(lambda: defaultdict(list))
# We also keep a flat map for the execution context: Name -> Class
# This handles potential name collisions by favoring the last seen or explicit logic if needed.
node_registry = {}
@ -33,14 +35,14 @@ def get_all_nodes():
for _, provider_name, ispkg in pkgutil.iter_modules(path, prefix):
if not ispkg:
continue
# e.g., provider_name = "diagrams.aws"
short_provider = provider_name.split(".")[-1]
# Skip internal modules if any (base, etc are actually useful, but we focus on providers)
if short_provider in ['base', 'custom']:
# 'custom' and 'base' might be treated differently, but for now we scan them
pass
if short_provider in ['base', 'custom']:
# 'custom' and 'base' might be treated differently, but for now we scan them
pass
try:
provider_module = importlib.import_module(provider_name)
@ -63,13 +65,14 @@ def get_all_nodes():
# or at least is defined in the diagrams package
if obj.__module__.startswith("diagrams"):
# Ensure the object actually belongs to this service (or a submodule of it)
# This prevents listing imported classes from other services (e.g. Trace in operations vs devtools)
# This prevents listing imported classes from other services (e.g. Trace
# in operations vs devtools)
if not obj.__module__.startswith(service_name):
continue
icons[short_provider][short_service].append(name)
node_registry[name] = obj
except ImportError:
continue
except Exception:

76
mcp-server/src/server.py
Unescape View File

@ -1,14 +1,15 @@
import base64
import contextlib
import os
import sys
import tempfile
import contextlib
import base64
from pathlib import Path
from mcp.server.fastmcp import FastMCP
from diagrams import Diagram, Cluster, Edge, Node
# Import our helper
from inspection import get_all_nodes
from mcp.server.fastmcp import FastMCP
from diagrams import Cluster, Diagram, Edge, Node
# Initialize FastMCP
mcp = FastMCP("diagrams-mcp")
@ -18,11 +19,12 @@ print("Loading diagram nodes...", file=sys.stderr)
ALL_ICONS, NODE_REGISTRY = get_all_nodes()
print(f"Loaded {len(NODE_REGISTRY)} nodes.", file=sys.stderr)
@mcp.tool()
def list_icons(provider_filter: str = None, service_filter: str = None):
"""
List available icons from the diagrams package, with optional filtering.
Args:
provider_filter: Filter icons by provider name (e.g., "aws", "gcp", "k8s")
service_filter: Filter icons by service name (e.g., "compute", "database")
@ -30,26 +32,30 @@ def list_icons(provider_filter: str = None, service_filter: str = None):
if not provider_filter:
# Return list of providers
return {"providers": list(ALL_ICONS.keys())}
if provider_filter not in ALL_ICONS:
return {"error": f"Provider '{provider_filter}' not found. Available: {list(ALL_ICONS.keys())}"}
provider_data = ALL_ICONS[provider_filter]
if not service_filter:
# Return all services for this provider
return provider_data
if service_filter not in provider_data:
return {"error": f"Service '{service_filter}' not found in '{provider_filter}'. Available: {list(provider_data.keys())}"}
return {
"error": f"Service '{service_filter}' not found in '{provider_filter}'. Available: {
list(
provider_data.keys())}"}
return {service_filter: provider_data[service_filter]}
@mcp.tool()
def get_diagram_examples(diagram_type: str = "all"):
"""
Get example code for different types of diagrams.
Args:
diagram_type: Type of diagram example to return (aws, k8s, flow, etc. or 'all')
"""
@ -70,11 +76,11 @@ from diagrams.k8s.network import Ingress, Service
with Diagram("K8s Cluster", show=False):
ingress = Ingress("domain.com")
with Cluster("App"):
svc = Service("svc")
pods = [Pod("pod1"), Pod("pod2")]
ingress >> svc >> pods
""",
"custom": """
@ -86,28 +92,29 @@ with Diagram("Custom", show=False):
Custom("Label", "./my-icon.png")
"""
}
if diagram_type == "all":
return examples
return {diagram_type: examples.get(diagram_type, "No example found for this type.")}
@mcp.tool()
def generate_diagram(code: str, filename: str = None, timeout: int = 90):
"""
Generate a diagram from Python code using the diagrams package.
Args:
code: Python code using the diagrams package DSL.
filename: Optional filename to save the diagram to.
timeout: Execution timeout in seconds.
"""
# Create a temporary directory for execution
with tempfile.TemporaryDirectory() as temp_dir:
original_cwd = os.getcwd()
os.chdir(temp_dir)
try:
# Prepare the execution context
# We inject Diagram, Cluster, Edge, and ALL discovered nodes (EC2, Pod, etc.)
@ -120,57 +127,60 @@ def generate_diagram(code: str, filename: str = None, timeout: int = 90):
"Node": Node,
**NODE_REGISTRY
}
# Execute the code
# We wrap it in a try/except block within the exec to catch runtime errors
try:
exec(code, exec_globals)
except Exception as e:
return {"status": "error", "message": f"Runtime error: {str(e)}"}
# Find the generated file
# Diagrams generates files based on the name passed to Diagram() class
# We look for any .png file created in the temp dir
generated_files = list(Path(".").glob("*.png"))
if not generated_files:
return {"status": "error", "message": "No diagram image was generated. Did you call with Diagram(..., show=False)?"}
return {
"status": "error",
"message": "No diagram image was generated. Did you call with Diagram(..., show=False)?"}
# Use the most recently modified file or the first one
generated_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
output_file = generated_files[0]
# If a filename was requested, we might want to rename it?
# For now, we return the path.
# For now, we return the path.
# In a real MCP setup, we might copy this to a mounted volume.
# Copy the generated file back to the original working directory
# This ensures that if the user mounted their project to the working directory,
# the file appears in their project.
import shutil
target_dir = Path(original_cwd)
target_filename = filename if filename else output_file.name
target_path = target_dir / target_filename
# Ensure extension
if not target_path.suffix:
target_path = target_path.with_suffix(".png")
shutil.copy2(output_file, target_path)
final_path = str(target_path)
return {
"status": "success",
"status": "success",
"path": final_path,
"filename": target_path.name
}
except Exception as e:
return {"status": "error", "message": f"System error: {str(e)}"}
finally:
os.chdir(original_cwd)
if __name__ == "__main__":
mcp.run()

Write Preview
Loading…
Cancel
Save