|
5 | 5 | import neo4j_viz as nvl |
6 | 6 | from pydantic import BaseModel, Field, ValidationInfo, field_validator |
7 | 7 |
|
| 8 | +NODE_COLOR_PALETTE = [ |
| 9 | + ("#e3f2fd", "#1976d2"), # Light Blue / Blue |
| 10 | + ("#f3e5f5", "#7b1fa2"), # Light Purple / Purple |
| 11 | + ("#e8f5e8", "#388e3c"), # Light Green / Green |
| 12 | + ("#fff3e0", "#f57c00"), # Light Orange / Orange |
| 13 | + ("#fce4ec", "#c2185b"), # Light Pink / Pink |
| 14 | + ("#e0f2f1", "#00695c"), # Light Teal / Teal |
| 15 | + ("#f1f8e9", "#689f38"), # Light Lime / Lime |
| 16 | + ("#fff8e1", "#ffa000"), # Light Amber / Amber |
| 17 | + ("#e8eaf6", "#3f51b5"), # Light Indigo / Indigo |
| 18 | + ("#efebe9", "#5d4037"), # Light Brown / Brown |
| 19 | + ("#fafafa", "#424242"), # Light Grey / Dark Grey |
| 20 | + ("#e1f5fe", "#0277bd"), # Light Cyan / Cyan |
| 21 | + ("#f9fbe7", "#827717"), # Light Yellow-Green / Olive |
| 22 | + ("#fff1f0", "#d32f2f"), # Light Red / Red |
| 23 | + ("#f4e6ff", "#6a1b9a"), # Light Violet / Violet |
| 24 | + ("#e6f7ff", "#1890ff"), # Very Light Blue / Bright Blue |
| 25 | +] |
| 26 | + |
8 | 27 |
|
9 | 28 | def _generate_relationship_pattern( |
10 | 29 | start_node_label: str, relationship_type: str, end_node_label: str |
@@ -147,6 +166,12 @@ def to_nvl(self) -> nvl.Node: |
147 | 166 | properties=self.all_properties_dict, |
148 | 167 | ) |
149 | 168 |
|
| 169 | + def get_mermaid_config_str(self) -> str: |
| 170 | + "Get the Mermaid configuration string for the node." |
| 171 | + props = [f"<br/>{self.key_property.name}: {self.key_property.type} | KEY"] |
| 172 | + props.extend([f"<br/>{p.name}: {p.type}" for p in self.properties]) |
| 173 | + return f'{self.label}["{self.label}{''.join(props)}"]' |
| 174 | + |
150 | 175 | @classmethod |
151 | 176 | def from_arrows(cls, arrows_node_dict: dict[str, Any]) -> "Node": |
152 | 177 | "Convert an Arrows Node to a Node." |
@@ -268,6 +293,16 @@ def to_nvl(self) -> nvl.Relationship: |
268 | 293 | properties=self.all_properties_dict, |
269 | 294 | ) |
270 | 295 |
|
| 296 | + def get_mermaid_config_str(self) -> str: |
| 297 | + "Get the Mermaid configuration string for the relationship." |
| 298 | + props = ( |
| 299 | + [f"<br/>{self.key_property.name}: {self.key_property.type} | KEY"] |
| 300 | + if self.key_property |
| 301 | + else [] |
| 302 | + ) |
| 303 | + props.extend([f"<br/>{p.name}: {p.type}" for p in self.properties]) |
| 304 | + return f"{self.start_node_label} -->|{self.type}{''.join(props)}| {self.end_node_label}" |
| 305 | + |
271 | 306 | @classmethod |
272 | 307 | def from_arrows( |
273 | 308 | cls, |
@@ -419,6 +454,33 @@ def to_nvl(self) -> nvl.VisualizationGraph: |
419 | 454 | relationships=[r.to_nvl() for r in self.relationships], |
420 | 455 | ) |
421 | 456 |
|
| 457 | + def _generate_mermaid_config_styling_str(self) -> str: |
| 458 | + "Generate the Mermaid configuration string for the data model." |
| 459 | + node_color_config = "" |
| 460 | + |
| 461 | + for idx, node in enumerate(self.nodes): |
| 462 | + node_color_config += f"classDef node_{idx}_color fill:{NODE_COLOR_PALETTE[idx % len(NODE_COLOR_PALETTE)][0]},stroke:{NODE_COLOR_PALETTE[idx % len(NODE_COLOR_PALETTE)][1]},stroke-width:3px,color:#000,font-size:12px\nclass {node.label} node_{idx}_color\n\n" |
| 463 | + |
| 464 | + return f""" |
| 465 | +%% Styling |
| 466 | +{node_color_config} |
| 467 | + """ |
| 468 | + |
| 469 | + def get_mermaid_config_str(self) -> str: |
| 470 | + "Get the Mermaid configuration string for the data model." |
| 471 | + mermaid_nodes = [n.get_mermaid_config_str() for n in self.nodes] |
| 472 | + mermaid_relationships = [r.get_mermaid_config_str() for r in self.relationships] |
| 473 | + mermaid_styling = self._generate_mermaid_config_styling_str() |
| 474 | + return f"""graph TD |
| 475 | +%% Nodes |
| 476 | +{"\n".join(mermaid_nodes)} |
| 477 | +
|
| 478 | +%% Relationships |
| 479 | +{"\n".join(mermaid_relationships)} |
| 480 | +
|
| 481 | +{mermaid_styling} |
| 482 | +""" |
| 483 | + |
422 | 484 | @classmethod |
423 | 485 | def from_arrows(cls, arrows_data_model_dict: dict[str, Any]) -> "DataModel": |
424 | 486 | "Convert an Arrows Data Model to a Data Model." |
|
0 commit comments