Start moving custom fields into properties dict

Only `Node`, `Relationship` and `from_neo4j` updated so far.

Co-Authored-By: Florentin Dörre <florentin.dorre@neotechnology.com>

Author: adamnsch

python-wrapper/src/neo4j_viz/neo4j.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
 import neo4j.graph
 from neo4j import Result
@@ -22,8 +21,8 @@ def from_neo4j(
     Create a VisualizationGraph from a Neo4j Graph or Neo4j Result object.
 
     All node and relationship properties will be included in the visualization graph.
-    If the property names are conflicting with those of `Node` and `Relationship` objects, they will be prefixed
-    with `__`.
+    If the properties are named as the fields of the `Node` or `Relationship` classes, they will be included as
+    top level fields of the respective objects. Otherwise, they will be included in the `properties` dictionary.
 
     Parameters
     ----------
@@ -63,69 +62,64 @@ def from_neo4j(
 
 
 def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_property: Optional[str]) -> Node:
-    labels = sorted([label for label in node.labels])
+    top_level_fields = {"id": node.element_id}
 
     if size_property:
-        size = node.get(size_property)
-    else:
-        size = None
+        top_level_fields["size"] = node.get(size_property)
 
+    labels = sorted([label for label in node.labels])
     if caption_property:
         if caption_property == "labels":
             if len(labels) > 0:
-                caption = ":".join([label for label in labels])
-            else:
-                caption = None
+                top_level_fields["caption"] = ":".join([label for label in labels])
         else:
-            caption = str(node.get(caption_property))
+            top_level_fields["caption"] = str(node.get(caption_property))
+
+    properties = {}
+    for prop, value in node.items():
+        if prop not in Node.model_fields.keys():
+            properties[prop] = value
+            continue
+
+        if prop in top_level_fields:
+            properties[prop] = value
+            continue
 
-    base_node_props = dict(id=node.element_id, caption=caption, labels=labels, size=size)
+        top_level_fields[prop] = value
 
-    protected_props = base_node_props.keys()
-    additional_node_props = {k: v for k, v in node.items()}
-    additional_node_props = _rename_protected_props(additional_node_props, protected_props)
+    if "labels" in properties:
+        properties["__labels"] = properties["labels"]
+    properties["labels"] = labels
 
-    return Node(**base_node_props, **additional_node_props)
+    return Node(**top_level_fields, properties=properties)
 
 
 def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[str]) -> Optional[Relationship]:
     if rel.start_node is None or rel.end_node is None:
         return None
 
+    top_level_fields = {"id": rel.element_id, "source": rel.start_node.element_id, "target": rel.end_node.element_id}
+
     if caption_property:
         if caption_property == "type":
-            caption = rel.type
+            top_level_fields["caption"] = rel.type
         else:
-            caption = str(rel.get(caption_property))
-    else:
-        caption = None
-
-    base_rel_props = dict(
-        id=rel.element_id,
-        source=rel.start_node.element_id,
-        target=rel.end_node.element_id,
-        _type=rel.type,
-        caption=caption,
-    )
-
-    protected_props = base_rel_props.keys()
-    additional_rel_props = {k: v for k, v in rel.items()}
-    additional_rel_props = _rename_protected_props(additional_rel_props, protected_props)
-
-    return Relationship(
-        **base_rel_props,
-        **additional_rel_props,
-    )
-
-
-def _rename_protected_props(
-    additional_props: dict[str, Any],
-    protected_props: Iterable[str],
-) -> dict[str, Union[str, int, float]]:
-    for prop in protected_props:
-        if prop not in additional_props:
+            top_level_fields["caption"] = str(rel.get(caption_property))
+
+    properties = {}
+    for prop, value in rel.items():
+        if prop not in Relationship.model_fields.keys():
+            properties[prop] = value
             continue
 
-        additional_props[f"__{prop}"] = additional_props.pop(prop)
+        if prop in top_level_fields:
+            properties[prop] = value
+            continue
+
+        top_level_fields[prop] = value
+
+    if "type" in properties:
+        properties["__type"] = properties["type"]
+    properties["type"] = rel.type
 
-    return additional_props
+    return Relationship(**top_level_fields, properties=properties)

python-wrapper/src/neo4j_viz/node.py

@@ -45,6 +45,8 @@ class Node(BaseModel, extra="allow"):
     x: Optional[RealNumber] = Field(None, description="The x-coordinate of the node")
     #: The y-coordinate of the node
     y: Optional[RealNumber] = Field(None, description="The y-coordinate of the node")
+    #: The properties of the node
+    properties: dict[str, Any] = Field(default_factory=dict, description="The properties of the node")
 
     @field_serializer("color")
     def serialize_color(self, color: Color) -> str:

python-wrapper/src/neo4j_viz/relationship.py

@@ -43,6 +43,8 @@ class Relationship(BaseModel, extra="allow"):
     )
     #: The color of the relationship. Allowed input is for example "#FF0000", "red" or (255, 0, 0)
     color: Optional[ColorType] = Field(None, description="The color of the relationship")
+    #: The properties of the relationship
+    properties: dict[str, Any] = Field(default_factory=dict, description="The properties of the relationship")
 
     @field_serializer("color")
     def serialize_color(self, color: Color) -> str:

python-wrapper/tests/test_neo4j.py

@@ -19,7 +19,7 @@ def graph_setup(neo4j_session: Session) -> Generator[None, None, None]:
 
 
 @pytest.mark.requires_neo4j_and_gds
-def test_from_neo4j_graph(neo4j_session: Session) -> None:
+def test_from_neo4j_graph_basic(neo4j_session: Session) -> None:
     graph = neo4j_session.run("MATCH (a:_CI_A|_CI_B)-[r]->(b) RETURN a, b, r ORDER BY a").graph()
 
     VG = from_neo4j(graph)
@@ -31,27 +31,31 @@ def test_from_neo4j_graph(neo4j_session: Session) -> None:
         Node(
             id=node_ids[0],
             caption="_CI_A",
-            labels=["_CI_A"],
-            name="Alice",
-            height=20,
-            __id=42,
-            _id=1337,
-            __caption="hello",
+            properties=dict(
+                labels=["_CI_A"],
+                name="Alice",
+                height=20,
+                id=42,
+                _id=1337,
+                caption="hello",
+            ),
         ),
         Node(
             id=node_ids[1],
             caption="_CI_A:_CI_B",
-            labels=["_CI_A", "_CI_B"],
-            name="Bob",
-            height=10,
-            __id=84,
-            __size=11,
-            __labels=[1, 2],
+            size=11,
+            properties=dict(
+                labels=["_CI_A", "_CI_B"],
+                name="Bob",
+                height=10,
+                id=84,
+                __labels=[1, 2],
+            ),
         ),
     ]
 
     assert len(VG.nodes) == 2
-    assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes  # type: ignore[attr-defined]
+    assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes
 
     assert len(VG.relationships) == 2
     vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")
@@ -76,27 +80,31 @@ def test_from_neo4j_result(neo4j_session: Session) -> None:
         Node(
             id=node_ids[0],
             caption="_CI_A",
-            labels=["_CI_A"],
-            name="Alice",
-            height=20,
-            __id=42,
-            _id=1337,
-            __caption="hello",
+            properties=dict(
+                labels=["_CI_A"],
+                name="Alice",
+                height=20,
+                id=42,
+                _id=1337,
+                caption="hello",
+            ),
         ),
         Node(
             id=node_ids[1],
             caption="_CI_A:_CI_B",
-            labels=["_CI_A", "_CI_B"],
-            name="Bob",
-            height=10,
-            __id=84,
-            __size=11,
-            __labels=[1, 2],
+            size=11,
+            properties=dict(
+                labels=["_CI_A", "_CI_B"],
+                name="Bob",
+                height=10,
+                id=84,
+                __labels=[1, 2],
+            ),
         ),
     ]
 
     assert len(VG.nodes) == 2
-    assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes  # type: ignore[attr-defined]
+    assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes
 
     assert len(VG.relationships) == 2
     vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")
@@ -119,29 +127,33 @@ def test_from_neo4j_graph_full(neo4j_session: Session) -> None:
         Node(
             id=node_ids[0],
             caption="Alice",
-            labels=["_CI_A"],
-            name="Alice",
-            height=20,
             size=60.0,
-            __id=42,
-            _id=1337,
-            __caption="hello",
+            properties=dict(
+                labels=["_CI_A"],
+                name="Alice",
+                height=20,
+                id=42,
+                _id=1337,
+                caption="hello",
+            ),
         ),
         Node(
             id=node_ids[1],
             caption="Bob",
-            labels=["_CI_A", "_CI_B"],
-            name="Bob",
-            height=10,
             size=3.0,
-            __id=84,
-            __size=11,
-            __labels=[1, 2],
+            properties=dict(
+                labels=["_CI_A", "_CI_B"],
+                name="Bob",
+                size=11,
+                height=10,
+                id=84,
+                __labels=[1, 2],
+            ),
         ),
     ]
 
     assert len(VG.nodes) == 2
-    assert sorted(VG.nodes, key=lambda x: x.name) == expected_nodes  # type: ignore[attr-defined]
+    assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes
 
     assert len(VG.relationships) == 2
     vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")

python-wrapper/tests/test_node.py

@@ -26,6 +26,7 @@ def test_nodes_with_all_options() -> None:
         "pinned": True,
         "x": 1,
         "y": 10,
+        "properties": {},
     }
 
 
@@ -36,6 +37,7 @@ def test_nodes_minimal_node() -> None:
 
     assert node.to_dict() == {
         "id": "1",
+        "properties": {},
     }
 
 
@@ -48,6 +50,7 @@ def test_node_with_float_size() -> None:
     assert node.to_dict() == {
         "id": "1",
         "size": 10.2,
+        "properties": {},
     }
 
 
@@ -60,6 +63,7 @@ def test_node_with_additional_fields() -> None:
     assert node.to_dict() == {
         "id": "1",
         "componentId": 2,
+        "properties": {},
     }
 
 
@@ -69,6 +73,7 @@ def test_id_aliases(alias: str) -> None:
 
     assert node.to_dict() == {
         "id": "1",
+        "properties": {},
     }
 
 

python-wrapper/tests/test_relationship.py

@@ -23,6 +23,7 @@ def test_rels_with_all_options() -> None:
         "captionAlign": "top",
         "captionSize": 12,
         "color": "#ff0000",
+        "properties": {},
     }
 
 
@@ -34,7 +35,7 @@ def test_rels_minimal_rel() -> None:
 
     rel_dict = rel.to_dict()
 
-    assert {"id", "from", "to"} == set(rel_dict.keys())
+    assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
     assert rel_dict["from"] == "1"
     assert rel_dict["to"] == "2"
 
@@ -43,12 +44,12 @@ def test_rels_additional_fields() -> None:
     rel = Relationship(
         source="1",
         target="2",
-        componentId=2,
+        properties=dict(componentId=2),
     )
 
     rel_dict = rel.to_dict()
-    assert {"id", "from", "to", "componentId"} == set(rel_dict.keys())
-    assert rel.componentId == 2  # type: ignore[attr-defined]
+    assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
+    assert rel.properties["componentId"] == 2
 
 
 @pytest.mark.parametrize("src_alias", ["source", "sourceNodeId", "source_node_id", "from"])
@@ -63,6 +64,6 @@ def test_aliases(src_alias: str, trg_alias: str) -> None:
 
     rel_dict = rel.to_dict()
 
-    assert {"id", "from", "to"} == set(rel_dict.keys())
+    assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
     assert rel_dict["from"] == "1"
     assert rel_dict["to"] == "2"