diff --git a/python-wrapper/src/neo4j_viz/node.py b/python-wrapper/src/neo4j_viz/node.py index 8c7cad0f..8f6c9580 100644 --- a/python-wrapper/src/neo4j_viz/node.py +++ b/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: diff --git a/python-wrapper/src/neo4j_viz/visualization_graph.py b/python-wrapper/src/neo4j_viz/visualization_graph.py index 373f912f..3233086b 100644 --- a/python-wrapper/src/neo4j_viz/visualization_graph.py +++ b/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, diff --git a/python-wrapper/tests/test_node.py b/python-wrapper/tests/test_node.py index 14ddc84f..25e70d23 100644 --- a/python-wrapper/tests/test_node.py +++ b/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, } diff --git a/python-wrapper/tests/test_pinned.py b/python-wrapper/tests/test_pinned.py new file mode 100644 index 00000000..dbec5f8d --- /dev/null +++ b/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 diff --git a/python-wrapper/tests/test_render.py b/python-wrapper/tests/test_render.py index 4989fffb..e823bcc7 100644 --- a/python-wrapper/tests/test_render.py +++ b/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"),