#!/usr/bin/env python3
"""
Simple Conection Graph Web UI with Graphviz integration.
This creates a web interface that:
1. Uses Graphviz for proper graph layouts
2. Serves focused views of connection graphs
"""
import os
import json
from http.server import SimpleHTTPRequestHandler
from urllib.parse import urlparse, parse_qs, unquote
from openmdao.core.conn_graph import GRAPH_COLORS
from openmdao.utils.file_utils import _load_and_exec
import openmdao.utils.hooks as hooks
[docs]
class ConnGraphHandler(SimpleHTTPRequestHandler):
"""
Custom handler for serving the connection graph web interface.
Parameters
----------
conn_graph : AllConnGraph
Connection graph instance used to serve UI requests.
*args : list
Positional arguments passed to the base handler.
**kwargs : dict
Keyword arguments passed to the base handler.
Attributes
----------
conn_graph : AllConnGraph
Connection graph instance used to serve UI requests.
"""
[docs]
def __init__(self, conn_graph, *args, **kwargs):
"""
Initialize the handler.
"""
self.conn_graph = conn_graph
super().__init__(*args, **kwargs)
[docs]
def do_GET(self):
"""
Handle GET requests.
"""
parsed_path = urlparse(self.path)
if parsed_path.path == '/':
self.serve_main_page()
elif parsed_path.path == '/api/graph_info':
self.serve_graph_info()
elif parsed_path.path.startswith('/api/subsystem/'):
subsystem = unquote(parsed_path.path.replace('/api/subsystem/', ''))
self.serve_subsystem_graph(subsystem)
elif parsed_path.path.startswith('/api/variable/'):
variable = unquote(parsed_path.path.replace('/api/variable/', ''))
self.serve_variable_graph(variable)
elif parsed_path.path == '/api/search':
self.serve_search()
else:
super().do_GET()
[docs]
def serve_main_page(self):
"""
Serve the main HTML page.
"""
html_content = self.get_html_template()
self.send_response(200)
self.send_header('Content-type', 'text/html')
self.end_headers()
self.wfile.write(html_content.encode())
[docs]
def serve_graph_info(self):
"""
Serve basic graph information.
"""
# Get all nodes with their data
nodes_data = {}
for node_id, node_data in self.conn_graph.nodes(data=True):
# Convert tuple node_id to JSON-serializable format
if isinstance(node_id, tuple) and len(node_id) == 2:
node_id_json = [node_id[0], node_id[1]]
else:
node_id_json = str(node_id)
nodes_data[json.dumps(node_id_json)] = {
'rel_name': node_data.get('rel_name', ''),
'pathname': node_data.get('pathname', ''),
'io': node_id[0] if isinstance(node_id, tuple) and len(node_id) == 2 else '',
'fillcolor': node_data.get('fillcolor', '')
}
info = {
'success': True,
'nodes': len(self.conn_graph.nodes()),
'edges': len(self.conn_graph.edges()),
'subsystems': self.get_subsystems(),
'nodes_data': nodes_data,
'graph_colors': GRAPH_COLORS
}
self.send_json_response(info)
[docs]
def serve_subsystem_graph(self, subsystem):
"""
Serve graph for a specific subsystem.
Parameters
----------
subsystem : str
Subsystem pathname to focus on.
"""
try:
# Get the DOT format string instead of pydot object
dot_string = self.conn_graph.get_dot(subsystem)
subgraph = self.conn_graph.get_drawable_graph(subsystem)
# get help colors from the subgraph (where fillcolor is set)
incolor = outcolor = None
for node_id, node_data in subgraph.nodes(data=True):
# Look for input nodes (starting with 'i')
if incolor is None and node_id[0] == 'i':
incolor = node_data.get('fillcolor', None)
# Look for output nodes (starting with 'o')
elif outcolor is None and node_id[0] == 'o':
outcolor = node_data.get('fillcolor', None)
# Stop once we have both colors
if incolor is not None and outcolor is not None:
break
help_colors = {'i': incolor, 'o': outcolor}
# Convert nodes to a format the frontend can use
nodes_data = {}
for node_id, node_data in subgraph.nodes(data=True):
# Extract the actual variable information from the node metadata
# The node_data contains the original rel_name and pathname
rel_name = node_data.get('rel_name', '')
pathname = node_data.get('pathname', '')
io_type = node_id[0]
# If io_type is not in node_data, extract from node_id tuple
if not io_type and isinstance(node_id, tuple) and len(node_id) == 2:
io_type = node_id[0] # 'i' or 'o'
nodes_data[str(node_id)] = {
'rel_name': rel_name,
'pathname': pathname,
'io': io_type,
'fillcolor': node_data.get('fillcolor', '')
}
response = {
'success': True,
'subsystem': subsystem,
'nodes': len(subgraph.nodes()),
'edges': len(subgraph.edges()),
'nodes_data': nodes_data,
'help_colors': help_colors,
'dot': dot_string,
}
except Exception as e:
print(f"Error in serve_subsystem_graph: {e}")
# Provide more helpful error message
error_msg = str(e)
if "not found" in error_msg.lower():
error_msg = f"Subsystem '{subsystem}' not found. Try searching for a different " \
"subsystem name."
response = {
'success': False,
'error': error_msg
}
self.send_json_response(response)
[docs]
def serve_variable_graph(self, variable):
"""
Serve graph focused on a specific variable.
Parameters
----------
variable : str
Variable name to focus on.
"""
try:
# Get the DOT format string using the connection graph method
dot_string = self.conn_graph.get_dot(varname=variable)
subgraph = self.conn_graph.get_drawable_graph(varname=variable)
# Count nodes and edges from the subgraph
nodes_count = len(subgraph.nodes())
edges_count = len(subgraph.edges())
response = {
'success': True,
'variable': variable,
'nodes': nodes_count,
'edges': edges_count,
'dot': dot_string
}
except Exception as e:
# Provide more helpful error message
error_msg = str(e)
if "not found" in error_msg.lower():
error_msg = f"Variable '{variable}' not found. Try searching for a different " \
"variable name."
response = {'success': False, 'error': error_msg}
self.send_json_response(response)
[docs]
def serve_search(self):
"""
Serve search results.
"""
query = self.get_query_param('q', '').lower()
if not query:
self.send_json_response({'results': []})
return
results = []
# Search nodes
for node_id, node_data in self.conn_graph.nodes(data=True):
# Try different name fields
name = node_data['rel_name']
if query in name.lower():
results.append({
'type': 'variable',
'name': name,
'path': node_data['pathname'],
'io_type': node_id[0]
})
# Search subsystems
subsystems = self.get_subsystems()
for subsystem in subsystems:
if query in subsystem.lower():
results.append({
'type': 'subsystem',
'name': subsystem,
'path': subsystem
})
self.send_json_response({'results': results[:20]})
[docs]
def get_subsystems(self):
"""
Get list of unique subsystems.
Returns
-------
list of str
Sorted list of subsystem pathnames.
"""
subsystems = {'model'}
for _, node_data in self.conn_graph.nodes(data=True):
pathname = node_data.get('pathname', '')
if pathname:
subsystems.add(pathname)
return sorted(subsystems)
[docs]
def create_text_graph(self, subgraph, title):
"""
Create a simple text-based graph representation when Graphviz is not available.
Parameters
----------
subgraph : networkx.DiGraph
Graph to represent.
title : str
Title for the generated view.
Returns
-------
str
HTML string containing a simple graph representation.
"""
html = '<div style="font-family: monospace; padding: 20px;">'
html += f'<h3>{title}</h3>'
html += f'<p><strong>Nodes:</strong> {len(subgraph.nodes())}</p>'
# Show nodes
html += '<h4>Variables:</h4><ul>'
for node_id, node_data in subgraph.nodes(data=True):
name = node_data['rel_name']
io_type = node_id[0]
io_label = 'Input' if io_type == 'i' else 'Output'
html += f'<li><strong>{name}</strong> ({io_label})</li>'
html += '</ul>'
# Show connections
html += '<h4>Connections:</h4><ul>'
for edge in subgraph.edges():
source_data = subgraph.nodes[edge[0]]['attrs']
target_data = subgraph.nodes[edge[1]]['attrs']
source_name = source_data['rel_name']
target_name = target_data['rel_name']
html += f'<li>{source_name} → {target_name}</li>'
html += '</ul>'
html += '<p><em>Note: Using text representation due to graph generation error</em></p>'
html += '</div>'
return html
[docs]
def get_query_param(self, param, default=''):
"""
Get query parameter value.
Parameters
----------
param : str
Query parameter name.
default : str
Default value if parameter is not present.
Returns
-------
str
Parameter value or the default.
"""
parsed_path = urlparse(self.path)
query_params = parse_qs(parsed_path.query)
return query_params.get(param, [default])[0]
[docs]
def send_json_response(self, data):
"""
Send JSON response.
Parameters
----------
data : dict
Object to serialize as JSON.
"""
self.send_response(200)
self.send_header('Content-type', 'application/json')
self.end_headers()
self.wfile.write(json.dumps(data).encode())
[docs]
def get_html_template(self):
"""
Get the HTML template.
Returns
-------
str
HTML template content.
"""
fpath = os.path.join(os.path.dirname(__file__), 'conn_graph_ui_template.html')
with open(fpath, "r", encoding="utf-8") as f:
html_content = f.read()
return html_content
def _conn_graph_setup_parser(parser):
"""
Set up the openmdao subparser for the 'openmdao graph' command.
Parameters
----------
parser : argparse subparser
The parser we're adding options to.
"""
parser.description = ('This command displays the connection graph for the specified variable '
'or system.')
parser.add_argument('file', nargs=1, help='Python file containing the model.')
parser.add_argument('--problem', action='store', dest='problem', help='Problem name')
parser.add_argument('-v', '--varname', action='store', dest='varname',
help='Show connection tree containing the given variable.')
parser.add_argument('--port', action='store', dest='port', help='Port number')
def _conn_graph_cmd(options, user_args):
"""
Execute the 'openmdao conn_graph' command.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
def _view_graph(model):
"""
Serve the connection graph UI.
Parameters
----------
model : System
Model owning the connection graph.
"""
if options.varname:
model.get_conn_graph().display(varname=options.varname)
else:
model.get_conn_graph().serve(port=options.port)
# register the hooks
def _set_dyn_hook(prob):
"""
Register setup hooks to display the graph.
Parameters
----------
prob : Problem
Problem instance being set up.
"""
hooks._register_hook('_setup_part2', class_name='Group', inst_id='',
post=_view_graph, exit=True)
hooks._setup_hooks(prob.model)
# register the hooks
hooks._register_hook('setup', 'Problem', pre=_set_dyn_hook, ncalls=1)
_load_and_exec(options.file[0], user_args)
