Add set_node_caption convenience method

changelog.md

@@ -2,15 +2,15 @@
 
 ## Breaking changes
 
-* Removed `table` property from nodes and relationships returned from `from_snowflake`, the table is represented by the `caption` field.
-* Changed default value of `override` parameter in `VisualizationGraph.color_nodes()` from `False` to `True`. The method now overrides existing node colors by default. To preserve existing colors, explicitly pass `override=False`.
+- Removed `table` property from nodes and relationships returned from `from_snowflake`, the table is represented by the `caption` field.
+- Changed default value of `override` parameter in `VisualizationGraph.color_nodes()` from `False` to `True`. The method now overrides existing node colors by default. To preserve existing colors, explicitly pass `override=False`.
 
 ## New features
 
+- Added `set_node_captions()` convenience method to `VisualizationGraph` for setting node captions from a field or property.
 
 ## Bug fixes
 
 ## Improvements
 
-
 ## Other changes

docs/source/customizing.rst

@@ -20,6 +20,45 @@ If you have not yet created a ``VisualizationGraph`` object, please refer to one
    :backlinks: none
 
 
+Setting node captions
+---------------------
+
+Node captions are the text labels displayed on nodes in the visualization.
+
+The ``set_node_captions`` method
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+By calling the :meth:`neo4j_viz.VisualizationGraph.set_node_captions` method, you can set node captions based on a
+node field (like ``id``, ``size``, etc.) or a node property (members of the ``Node.properties`` map).
+
+The method accepts an ``override`` parameter (default ``True``) that controls whether to replace existing captions.
+If ``override=False``, only nodes without captions will be updated.
+
+Here's an example of setting node captions from a property:
+
+.. code-block:: python
+
+    # VG is a VisualizationGraph object with nodes that have a "name" property
+    VG.set_node_captions(property="name")
+
+You can also set captions from a node field, and choose not to override existing captions:
+
+.. code-block:: python
+
+    # VG is a VisualizationGraph object
+    VG.set_node_captions(field="id", override=False)
+
+For more complex scenarios where you need fallback logic or want to combine multiple properties, you can iterate over
+nodes directly:
+
+.. code-block:: python
+
+    # VG is a VisualizationGraph object
+    for node in VG.nodes:
+        caption = node.properties.get("name") or node.properties.get("title") or node.id
+        node.caption = str(caption)
+
+
 Coloring nodes
 --------------
 

python-wrapper/src/neo4j_viz/visualization_graph.py

@@ -194,6 +194,75 @@ def toggle_nodes_pinned(self, pinned: dict[NodeIdType, bool]) -> None:
 
             node.pinned = node_pinned
 
+    def set_node_captions(
+        self,
+        *,
+        field: Optional[str] = None,
+        property: Optional[str] = None,
+        override: bool = True,
+    ) -> None:
+        """
+        Set the caption for nodes in the graph based on either a node field or a node property.
+
+        Parameters
+        ----------
+        field:
+            The field of the nodes to use as the caption. Must be None if `property` is provided.
+        property:
+            The property of the nodes to use as the caption. Must be None if `field` is provided.
+        override:
+            Whether to override existing captions of the nodes, if they have any.
+
+        Examples
+        --------
+        Given a VisualizationGraph `VG`:
+
+        >>> nodes = [
+        ...    Node(id="0", properties={"name": "Alice", "age": 30}),
+        ...    Node(id="1", properties={"name": "Bob", "age": 25}),
+        ... ]
+        >>> VG = VisualizationGraph(nodes=nodes)
+
+        Set node captions from a property:
+
+        >>> VG.set_node_captions(property="name")
+
+        Set node captions from a field, only if not already set:
+
+        >>> VG.set_node_captions(field="id", override=False)
+
+        Set captions from multiple properties with fallback:
+
+        >>> for node in VG.nodes:
+        ...     caption = node.properties.get("name") or node.properties.get("title") or node.id
+        ...     if override or node.caption is None:
+        ...         node.caption = str(caption)
+        """
+        if not ((field is None) ^ (property is None)):
+            raise ValueError(
+                f"Exactly one of the arguments `field` (received '{field}') and `property` (received '{property}') must be provided"
+            )
+
+        if property:
+            # Use property
+            for node in self.nodes:
+                if not override and node.caption is not None:
+                    continue
+
+                value = node.properties.get(property, "")
+                node.caption = str(value)
+        else:
+            # Use field
+            assert field is not None
+            attribute = to_snake(field)
+
+            for node in self.nodes:
+                if not override and node.caption is not None:
+                    continue
+
+                value = getattr(node, attribute, "")
+                node.caption = str(value)
+
     def resize_nodes(
         self,
         sizes: Optional[dict[NodeIdType, RealNumber]] = None,

python-wrapper/tests/test_captions.py

@@ -0,0 +1,161 @@
+import pytest
+
+from neo4j_viz import Node, VisualizationGraph
+
+
+def test_set_node_captions_from_property() -> None:
+    """Test setting captions from a node property."""
+    nodes = [
+        Node(id="1", properties={"name": "Alice", "age": 30}),
+        Node(id="2", properties={"name": "Bob", "age": 25}),
+        Node(id="3", properties={"name": "Charlie", "age": 35}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="name")
+
+    assert VG.nodes[0].caption == "Alice"
+    assert VG.nodes[1].caption == "Bob"
+    assert VG.nodes[2].caption == "Charlie"
+
+
+def test_set_node_captions_from_field() -> None:
+    """Test setting captions from a node field."""
+    nodes = [
+        Node(id="node-1", properties={"name": "Alice"}),
+        Node(id="node-2", properties={"name": "Bob"}),
+        Node(id="node-3", properties={"name": "Charlie"}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(field="id")
+
+    assert VG.nodes[0].caption == "node-1"
+    assert VG.nodes[1].caption == "node-2"
+    assert VG.nodes[2].caption == "node-3"
+
+
+def test_set_node_captions_override_true() -> None:
+    """Test that override=True replaces existing captions."""
+    nodes = [
+        Node(id="1", caption="OldCaption1", properties={"name": "Alice"}),
+        Node(id="2", caption="OldCaption2", properties={"name": "Bob"}),
+        Node(id="3", properties={"name": "Charlie"}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="name", override=True)
+
+    assert VG.nodes[0].caption == "Alice"
+    assert VG.nodes[1].caption == "Bob"
+    assert VG.nodes[2].caption == "Charlie"
+
+
+def test_set_node_captions_override_false() -> None:
+    """Test that override=False preserves existing captions."""
+    nodes = [
+        Node(id="1", caption="ExistingCaption", properties={"name": "Alice"}),
+        Node(id="2", properties={"name": "Bob"}),
+        Node(id="3", caption="AnotherCaption", properties={"name": "Charlie"}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="name", override=False)
+
+    assert VG.nodes[0].caption == "ExistingCaption"  # Not overridden
+    assert VG.nodes[1].caption == "Bob"  # Set (was None)
+    assert VG.nodes[2].caption == "AnotherCaption"  # Not overridden
+
+
+def test_set_node_captions_missing_property() -> None:
+    """Test behavior when property is missing from some nodes."""
+    nodes = [
+        Node(id="1", properties={"name": "Alice"}),
+        Node(id="2", properties={"age": 25}),  # No "name" property
+        Node(id="3", properties={"name": "Charlie"}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="name")
+
+    assert VG.nodes[0].caption == "Alice"
+    assert VG.nodes[1].caption == ""  # Empty string for missing property
+    assert VG.nodes[2].caption == "Charlie"
+
+
+def test_set_node_captions_numeric_property() -> None:
+    """Test setting captions from numeric properties."""
+    nodes = [
+        Node(id="1", properties={"score": 100}),
+        Node(id="2", properties={"score": 200}),
+        Node(id="3", properties={"score": 300}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="score")
+
+    assert VG.nodes[0].caption == "100"
+    assert VG.nodes[1].caption == "200"
+    assert VG.nodes[2].caption == "300"
+
+
+def test_set_node_captions_field_and_property_both_provided() -> None:
+    """Test that providing both field and property raises an error."""
+    nodes = [Node(id="1", properties={"name": "Alice"})]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    with pytest.raises(ValueError, match="Exactly one of the arguments"):
+        VG.set_node_captions(field="id", property="name")
+
+
+def test_set_node_captions_neither_field_nor_property() -> None:
+    """Test that providing neither field nor property raises an error."""
+    nodes = [Node(id="1", properties={"name": "Alice"})]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    with pytest.raises(ValueError, match="Exactly one of the arguments"):
+        VG.set_node_captions()
+
+
+def test_set_node_captions_field_with_snake_case() -> None:
+    """Test that field names are converted to snake_case."""
+    nodes = [
+        Node(id="1", caption_size=1, properties={"name": "Alice"}),
+        Node(id="2", caption_size=2, properties={"name": "Bob"}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(field="captionSize")
+
+    assert VG.nodes[0].caption == "1"
+    assert VG.nodes[1].caption == "2"
+
+
+def test_set_node_captions_empty_graph() -> None:
+    """Test setting captions on an empty graph."""
+    VG = VisualizationGraph(nodes=[], relationships=[])
+
+    # Should not raise an error
+    VG.set_node_captions(property="name")
+
+    assert len(VG.nodes) == 0
+
+
+def test_set_node_captions_complex_property_values() -> None:
+    """Test setting captions from properties with complex types."""
+    nodes = [
+        Node(id="1", properties={"tags": ["tag1", "tag2"]}),
+        Node(id="2", properties={"metadata": {"key": "value"}}),
+        Node(id="3", properties={"value": None}),
+    ]
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    VG.set_node_captions(property="tags")
+    assert VG.nodes[0].caption == "['tag1', 'tag2']"
+    assert VG.nodes[1].caption == ""
+    assert VG.nodes[2].caption == ""
+
+    VG.set_node_captions(property="metadata")
+    assert VG.nodes[0].caption == ""
+    assert VG.nodes[1].caption == "{'key': 'value'}"
+    assert VG.nodes[2].caption == ""