Add more params to from_gql_create

Author: adamnsch

python-wrapper/src/neo4j_viz/gql_create.py

@@ -90,7 +90,13 @@ def _get_snippet(q: str, idx: int, context: int = 15) -> str:
     return q[start:end].replace("\n", " ")
 
 
-def from_gql_create(query: str) -> VisualizationGraph:
+def from_gql_create(
+    query: str,
+    size_property: Optional[str] = None,
+    node_caption: Optional[str] = "labels",
+    relationship_caption: Optional[str] = "type",
+    node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+) -> VisualizationGraph:
     """
     Parse a GQL CREATE query and return a VisualizationGraph object representing the graph it creates.
 
@@ -107,6 +113,15 @@ def from_gql_create(query: str) -> VisualizationGraph:
     ----------
     query : str
         The GQL CREATE query to parse
+    size_property : str, optional
+        Property to use for node size, by default None.
+    node_caption : str, optional
+        Property to use as the node caption, by default the node labels will be used.
+    relationship_caption : str, optional
+        Property to use as the relationship caption, by default the relationship type will be used.
+    node_radius_min_max : tuple[float, float], optional
+        Minimum and maximum node radius, by default (3, 60).
+        To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
     """
 
     query = query.strip()
@@ -295,4 +310,28 @@ def parse_labels_and_props(
                 snippet = part[:30]
                 raise ValueError(f"Invalid element in CREATE near: `{snippet}`.")
 
-    return VisualizationGraph(nodes=nodes, relationships=relationships)
+    if size_property is not None:
+        for node in nodes:
+            node.size = node.properties.get(size_property)
+
+    if node_caption is not None:
+        for node in nodes:
+            if node_caption == "labels":
+                if len(node.properties["labels"]) > 0:
+                    node.caption = ":".join([label for label in node.properties["labels"]])
+            else:
+                node.caption = str(node.properties.get(node_caption))
+
+    if relationship_caption is not None:
+        for rel in relationships:
+            if relationship_caption == "type":
+                rel.caption = rel.properties["type"]
+            else:
+                rel.caption = str(rel.properties.get(relationship_caption))
+
+    VG = VisualizationGraph(nodes=nodes, relationships=relationships)
+
+    if (node_radius_min_max is not None) and (size_property is not None):
+        VG.resize_nodes(node_radius_min_max=node_radius_min_max)
+
+    return VG

python-wrapper/src/neo4j_viz/neo4j.py

@@ -56,7 +56,7 @@ def from_neo4j(
 
     VG = VisualizationGraph(nodes, relationships)
 
-    if node_radius_min_max and size_property is not None:
+    if (node_radius_min_max is not None) and (size_property is not None):
         VG.resize_nodes(node_radius_min_max=node_radius_min_max)
 
     return VG

python-wrapper/tests/test_gql_create.py

@@ -5,7 +5,7 @@
 from neo4j_viz.gql_create import from_gql_create
 
 
-def test_from_gql_create() -> None:
+def test_from_gql_create_syntax() -> None:
     query = """
             CREATE
               (a:User {name: 'Alice', age: 23, labels: ['Happy'], "id": 42}),
@@ -47,7 +47,7 @@ def test_from_gql_create() -> None:
         {"top_level": {}, "properties": {"name": "Fawad", "age": 78, "labels": ["Person", "User"]}},
     ]
 
-    VG = from_gql_create(query)
+    VG = from_gql_create(query, node_caption=None, relationship_caption=None)
 
     assert len(VG.nodes) == len(expected_node_dicts)
     for i, exp_node in enumerate(expected_node_dicts):
@@ -80,6 +80,81 @@ def test_from_gql_create() -> None:
         assert created_rel.properties == exp_rel["properties"]
 
 
+def test_from_gql_create_captions() -> None:
+    query = """
+            CREATE
+              (a:User {name: 'Alice', age: 23}),
+              (b:User:person {name: "Bridget", age: 34, "caption": "Bridget"}),
+              (a)-[:LINK {weight: 0.5}]->(b);
+            """
+    expected_node_dicts: list[dict[str, dict[str, Any]]] = [
+        {
+            "top_level": {"caption": "User"},
+            "properties": {"name": "Alice", "age": 23, "labels": ["User"]},
+        },
+        {
+            "top_level": {"caption": "User:person"},
+            "properties": {"name": "Bridget", "age": 34, "labels": ["User", "person"]},
+        },
+    ]
+
+    VG = from_gql_create(query)
+
+    assert len(VG.nodes) == len(expected_node_dicts)
+    for i, exp_node in enumerate(expected_node_dicts):
+        created_node = VG.nodes[i]
+
+        assert created_node.model_dump(exclude_none=True, exclude={"properties", "id"}) == exp_node["top_level"]
+        assert created_node.properties == exp_node["properties"]
+
+    expected_relationships_dicts: list[dict[str, Any]] = [
+        {
+            "source_idx": 0,
+            "target_idx": 1,
+            "top_level": {"caption": "LINK"},
+            "properties": {"weight": 0.5, "type": "LINK"},
+        },
+    ]
+
+    assert len(VG.relationships) == len(expected_relationships_dicts)
+    for i, exp_rel in enumerate(expected_relationships_dicts):
+        created_rel = VG.relationships[i]
+        assert created_rel.source == VG.nodes[exp_rel["source_idx"]].id
+        assert created_rel.target == VG.nodes[exp_rel["target_idx"]].id
+        assert (
+            created_rel.model_dump(exclude_none=True, exclude={"properties", "id", "source", "target"})
+            == exp_rel["top_level"]
+        )
+        assert created_rel.properties == exp_rel["properties"]
+
+
+def test_from_gql_create_sizes() -> None:
+    query = """
+            CREATE
+              (a:User {name: 'Alice', age: 23}),
+              (b:User:person {name: "Bridget", age: 34, "caption": "Bridget"});
+            """
+    expected_node_dicts: list[dict[str, dict[str, Any]]] = [
+        {
+            "top_level": {"size": 3.0},
+            "properties": {"name": "Alice", "age": 23, "labels": ["User"]},
+        },
+        {
+            "top_level": {"caption": "Bridget", "size": 60.0},
+            "properties": {"name": "Bridget", "age": 34, "labels": ["User", "person"]},
+        },
+    ]
+
+    VG = from_gql_create(query, size_property="age", node_caption=None, relationship_caption=None)
+
+    assert len(VG.nodes) == len(expected_node_dicts)
+    for i, exp_node in enumerate(expected_node_dicts):
+        created_node = VG.nodes[i]
+
+        assert created_node.model_dump(exclude_none=True, exclude={"properties", "id"}) == exp_node["top_level"]
+        assert created_node.properties == exp_node["properties"]
+
+
 def test_unbalanced_parentheses_snippet() -> None:
     query = "CREATE (a:User, (b:User })"
     with pytest.raises(ValueError, match=r"Unbalanced parentheses near: `.*\(b:User.*"):