chore: fix pre-commit hook failures and formatting

pull/1203/head
Carlos Barbero 8 months ago
parent 57433dc2e2
commit adb8661cb7

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

@ -1,11 +1,13 @@
import pkgutil
import importlib import importlib
import inspect import inspect
import pkgutil
import sys import sys
from collections import defaultdict from collections import defaultdict
import diagrams import diagrams
from diagrams import Node from diagrams import Node
def get_all_nodes(): def get_all_nodes():
""" """
Dynamically inspects the diagrams package and returns a dictionary of all available Nodes. Dynamically inspects the diagrams package and returns a dictionary of all available Nodes.
@ -63,7 +65,8 @@ def get_all_nodes():
# or at least is defined in the diagrams package # or at least is defined in the diagrams package
if obj.__module__.startswith("diagrams"): if obj.__module__.startswith("diagrams"):
# Ensure the object actually belongs to this service (or a submodule of it) # 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): if not obj.__module__.startswith(service_name):
continue continue

@ -1,14 +1,15 @@
import base64
import contextlib
import os import os
import sys import sys
import tempfile import tempfile
import contextlib
import base64
from pathlib import Path from pathlib import Path
from mcp.server.fastmcp import FastMCP
from diagrams import Diagram, Cluster, Edge, Node
# Import our helper # Import our helper
from inspection import get_all_nodes from inspection import get_all_nodes
from mcp.server.fastmcp import FastMCP
from diagrams import Cluster, Diagram, Edge, Node
# Initialize FastMCP # Initialize FastMCP
mcp = FastMCP("diagrams-mcp") mcp = FastMCP("diagrams-mcp")
@ -18,6 +19,7 @@ print("Loading diagram nodes...", file=sys.stderr)
ALL_ICONS, NODE_REGISTRY = get_all_nodes() ALL_ICONS, NODE_REGISTRY = get_all_nodes()
print(f"Loaded {len(NODE_REGISTRY)} nodes.", file=sys.stderr) print(f"Loaded {len(NODE_REGISTRY)} nodes.", file=sys.stderr)
@mcp.tool() @mcp.tool()
def list_icons(provider_filter: str = None, service_filter: str = None): def list_icons(provider_filter: str = None, service_filter: str = None):
""" """
@ -41,10 +43,14 @@ def list_icons(provider_filter: str = None, service_filter: str = None):
return provider_data return provider_data
if service_filter not in 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]} return {service_filter: provider_data[service_filter]}
@mcp.tool() @mcp.tool()
def get_diagram_examples(diagram_type: str = "all"): def get_diagram_examples(diagram_type: str = "all"):
""" """
@ -92,6 +98,7 @@ with Diagram("Custom", show=False):
return {diagram_type: examples.get(diagram_type, "No example found for this type.")} return {diagram_type: examples.get(diagram_type, "No example found for this type.")}
@mcp.tool() @mcp.tool()
def generate_diagram(code: str, filename: str = None, timeout: int = 90): def generate_diagram(code: str, filename: str = None, timeout: int = 90):
""" """
@ -134,7 +141,9 @@ def generate_diagram(code: str, filename: str = None, timeout: int = 90):
generated_files = list(Path(".").glob("*.png")) generated_files = list(Path(".").glob("*.png"))
if not generated_files: 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 # Use the most recently modified file or the first one
generated_files.sort(key=lambda f: f.stat().st_mtime, reverse=True) generated_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
@ -172,5 +181,6 @@ def generate_diagram(code: str, filename: str = None, timeout: int = 90):
finally: finally:
os.chdir(original_cwd) os.chdir(original_cwd)
if __name__ == "__main__": if __name__ == "__main__":
mcp.run() mcp.run()

Loading…
Cancel
Save