Created
February 15, 2025 01:47
-
-
Save z-ai-lab/d0e669fb4e89331dee532311dee0b362 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import json | |
| import networkx as nx | |
| from collections import defaultdict | |
| from typing import Dict, Set, List | |
| def export_schema_to_arrows(G: nx.DiGraph) -> str: | |
| """ | |
| Export the schema of a NetworkX DiGraph to Arrows.app JSON format. | |
| Schema includes unique node labels and relationship types with their connections. | |
| Args: | |
| G (nx.DiGraph): NetworkX directed graph | |
| Returns: | |
| str: JSON string compatible with Arrows.app import | |
| """ | |
| # Collect unique node labels and their relationships | |
| node_labels: Set[str] = set() | |
| relationships: Dict[str, Set[tuple]] = defaultdict(set) | |
| # Extract node labels | |
| for _, node_data in G.nodes(data=True): | |
| label = node_data.get('label', 'Node') | |
| node_labels.add(label) | |
| # Extract relationship types and their connections | |
| for source, target, edge_data in G.edges(data=True): | |
| source_label = G.nodes[source].get('label', 'Node') | |
| target_label = G.nodes[target].get('label', 'Node') | |
| rel_type = edge_data.get('type', 'RELATED_TO') | |
| relationships[rel_type].add((source_label, target_label)) | |
| # Initialize arrows.app structure | |
| arrows_data = { | |
| "nodes": [], | |
| "relationships": [], | |
| "style": { | |
| "font-family": "sans-serif", | |
| "background-color": "#ffffff", | |
| "background-image": "", | |
| "background-size": "100%", | |
| "node-color": "#ffffff", | |
| "border-width": 4, | |
| "border-color": "#000000", | |
| "radius": 50, | |
| "node-padding": 5, | |
| "node-margin": 2, | |
| "outside-position": "auto", | |
| "node-icon-image": "", | |
| "node-background-image": "", | |
| "icon-position": "inside", | |
| "icon-size": 64, | |
| "caption-position": "inside", | |
| "caption-max-width": 200, | |
| "caption-color": "#000000", | |
| "caption-font-size": 50, | |
| "caption-font-weight": "normal", | |
| "label-position": "outside", | |
| "label-display": "pill", | |
| "label-color": "#000000", | |
| "label-background-color": "#ffffff", | |
| "label-border-style": "solid", | |
| "label-border-width": 4, | |
| "label-border-color": "#000000", | |
| "label-font-size": 40, | |
| "label-padding": 5, | |
| "label-margin": 4, | |
| "directionality": "directed", | |
| "detail-position": "inline", | |
| "detail-orientation": "parallel", | |
| "arrow-width": 5, | |
| "arrow-color": "#000000", | |
| "margin-start": 5, | |
| "margin-end": 5, | |
| "margin-peer": 20, | |
| "attachment-start": "normal", | |
| "attachment-end": "normal", | |
| "relationship-icon-image": "", | |
| "type-color": "#000000", | |
| "type-background-color": "#ffffff", | |
| "type-border-style": "none", | |
| "type-border-width": 0, | |
| "type-border-color": "#000000", | |
| "type-font-size": 16, | |
| "type-padding": 5, | |
| "property-position": "outside", | |
| "property-alignment": "colon", | |
| "property-color": "#000000", | |
| "property-font-size": 16, | |
| "property-font-weight": "normal" | |
| } | |
| } | |
| # Add nodes to schema diagram | |
| node_id_map = {} # Maps label to node ID | |
| for i, label in enumerate(sorted(node_labels)): | |
| node_id = f"n{i}" | |
| node_id_map[label] = node_id | |
| # Position nodes in a circle | |
| angle = (2 * 3.14159 * i) / len(node_labels) | |
| radius = 200 | |
| x = 300 + radius * math.cos(angle) | |
| y = 300 + radius * math.sin(angle) | |
| node = { | |
| "id": node_id, | |
| "position": { | |
| "x": x, | |
| "y": y | |
| }, | |
| "caption": label, | |
| "labels": [label], | |
| "properties": {} | |
| } | |
| arrows_data['nodes'].append(node) | |
| # Add relationships to schema diagram | |
| rel_id = 0 | |
| added_relationships = set() # Track unique relationship patterns | |
| for rel_type, connections in relationships.items(): | |
| for source_label, target_label in connections: | |
| # Create a unique key for this relationship pattern | |
| rel_key = (source_label, rel_type, target_label) | |
| # Only add if we haven't seen this pattern before | |
| if rel_key not in added_relationships: | |
| relationship = { | |
| "id": f"r{rel_id}", | |
| "type": rel_type, | |
| "fromId": node_id_map[source_label], | |
| "toId": node_id_map[target_label], | |
| "properties": {} | |
| } | |
| arrows_data['relationships'].append(relationship) | |
| added_relationships.add(rel_key) | |
| rel_id += 1 | |
| return json.dumps(arrows_data, indent=2) | |
| # Example usage: | |
| if __name__ == "__main__": | |
| import math # Required for node positioning | |
| # Create a sample graph | |
| G = nx.DiGraph() | |
| # Add nodes with labels | |
| G.add_node(1, label="Person") | |
| G.add_node(2, label="Person") | |
| G.add_node(3, label="Company") | |
| G.add_node(4, label="Product") | |
| # Add edges with types | |
| G.add_edge(1, 2, type="KNOWS") | |
| G.add_edge(1, 3, type="WORKS_AT") | |
| G.add_edge(3, 4, type="PRODUCES") | |
| G.add_edge(1, 4, type="BOUGHT") | |
| # Export schema to Arrows.app format | |
| schema_json = export_schema_to_arrows(G) | |
| print(schema_json) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment