Merge pull request #59 from neo4j/pinning

Add node pinning support

Author: FlorentinD

python-wrapper/src/neo4j_viz/node.py

@@ -33,6 +33,7 @@ class Node(BaseModel, extra="allow"):
     )
     size: Optional[RealNumber] = Field(None, ge=0, description="The size of the node as radius in pixel")
     color: Optional[ColorType] = Field(None, description="The color of the node")
+    pinned: Optional[bool] = Field(None, description="Whether the node is pinned in the visualization")
 
     @field_serializer("color")
     def serialize_color(self, color: Color) -> str:

python-wrapper/src/neo4j_viz/visualization_graph.py

@@ -66,6 +66,23 @@ def render(
             height,
         )
 
+    def toggle_nodes_pinned(self, pinned: dict[NodeIdType, bool]) -> None:
+        """
+        Toggle whether nodes should be pinned or not.
+
+        Parameters
+        ----------
+        pinned:
+            A dictionary mapping from node ID to whether the node should be pinned or not.
+        """
+        for node in self.nodes:
+            node_pinned = pinned.get(node.id)
+
+            if node_pinned is None:
+                continue
+
+            node.pinned = node_pinned
+
     def resize_nodes(
         self,
         sizes: Optional[dict[NodeIdType, RealNumber]] = None,

python-wrapper/tests/test_node.py

@@ -11,6 +11,7 @@ def test_nodes_with_all_options() -> None:
         caption_size=1,
         color="#FF0000",
         size=10,
+        pinned=True,
     )
 
     assert node.to_dict() == {
@@ -20,6 +21,7 @@ def test_nodes_with_all_options() -> None:
         "captionSize": 1,
         "color": "#ff0000",
         "size": 10,
+        "pinned": True,
     }
 
 

python-wrapper/tests/test_pinned.py

@@ -0,0 +1,20 @@
+from neo4j_viz import Node, VisualizationGraph
+from neo4j_viz.node import NodeIdType
+
+
+def test_toggle_nodes_pinned() -> None:
+    nodes = [
+        Node(id=42, pinned=False),
+        Node(id=43),
+        Node(id=44),
+        Node(id="1337", pinned=True),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    pinned: dict[NodeIdType, bool] = {42: True, 43: True, "1337": False}
+    VG.toggle_nodes_pinned(pinned)
+
+    assert VG.nodes[0].pinned
+    assert VG.nodes[1].pinned
+    assert VG.nodes[2].pinned is None
+    assert not VG.nodes[3].pinned

python-wrapper/tests/test_render.py

@@ -20,7 +20,7 @@ def test_basic_render(render_option: dict[str, Any], tmp_path: Path) -> None:
     nodes = [
         Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:0", caption="Person"),
         Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:6", caption="Product"),
-        Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:11", caption="Product"),
+        Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:11", caption="Product", pinned=True),
         Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:12", caption="Product"),
         Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:1", caption="Person"),
         Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:7", caption="Product"),