Merge pull request #176 from neo4j/expose-layout-options

expose layout options

Author: FlorentinD

changelog.md

@@ -7,7 +7,9 @@
 ## New features
 
 * Allow visualization based only on relationship DataFrames, without specifying node DataFrames in `from_dfs`
-* Add relationship properties to `VisualizationGraph` when constructing via `from_gds`.
+* Add relationship properties to `VisualizationGraph` when constructing via `from_gds`
+* Allow setting `layout_options` for `VisualizationGraph::render`
+
 
 ## Bug fixes
 

docs/source/api-reference/render_options.rst

@@ -4,5 +4,20 @@
 .. autoenum:: neo4j_viz.Layout
     :members:
 
+.. autoclass:: neo4j_viz.Force-Directed
+    :members:
+    :exclude-members: model_config
+
+.. autoclass:: neo4j_viz.HierarchicalLayoutOptions
+    :members:
+    :exclude-members: model_config
+
+.. autoenum:: neo4j_viz.options.Direction
+    :members:
+
+.. autoenum:: neo4j_viz.options.Packing
+    :members:
+
+
 .. autoenum:: neo4j_viz.Renderer
     :members:

docs/source/rendering.rst

@@ -35,8 +35,8 @@ You can provide these either as a percentage of the available space (eg. ``"80%"
 
 Further you can change the layout of the graph using the ``layout`` parameter, which can be set to one of the following values:
 
-* ``Layout.FORCE_DIRECTED`` - Nodes are arranged using the Force-Directed algorithm, which simulates physical forces
-* ``Layout.HIERARCHICAL`` - Arranges nodes by the directionality of their relationships, creating a tree-like structure
+* ``Layout.FORCE_DIRECTED`` - Nodes are arranged using the Force-Directed algorithm, which simulates physical forces. To customize the layout, use `ForceDirectedOptions` via `layout_options`.`
+* ``Layout.HIERARCHICAL`` - Arranges nodes by the directionality of their relationships, creating a tree-like structure.  To customize the layout use `HierarchicalLayoutOptions` via `layout_options`.`
 * ``Layout.COORDINATE`` - Arranges nodes based on coordinates defined in `x` and `y` properties on each node.
 
 Another thing of note is the ``max_allowed_nodes`` parameter, which controls the maximum number of nodes that is allowed

python-wrapper/src/neo4j_viz/__init__.py

@@ -1,6 +1,25 @@
 from .node import Node
-from .options import CaptionAlignment, Layout, Renderer
+from .options import (
+    CaptionAlignment,
+    Direction,
+    ForceDirectedLayoutOptions,
+    HierarchicalLayoutOptions,
+    Layout,
+    Packing,
+    Renderer,
+)
 from .relationship import Relationship
 from .visualization_graph import VisualizationGraph
 
-__all__ = ["VisualizationGraph", "Node", "Relationship", "CaptionAlignment", "Layout", "Renderer"]
+__all__ = [
+    "VisualizationGraph",
+    "Node",
+    "Relationship",
+    "CaptionAlignment",
+    "Layout",
+    "Renderer",
+    "ForceDirectedLayoutOptions",
+    "HierarchicalLayoutOptions",
+    "Direction",
+    "Packing",
+]

python-wrapper/src/neo4j_viz/options.py

@@ -2,10 +2,10 @@
 
 import warnings
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import enum_tools.documentation
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError, model_validator
 
 
 @enum_tools.documentation.document_enum
@@ -33,6 +33,69 @@ class Layout(str, Enum):
     GRID = "grid"
 
 
+@enum_tools.documentation.document_enum
+class Direction(str, Enum):
+    """
+    The direction in which the layout should be oriented
+    """
+
+    LEFT = "left"
+    RIGHT = "right"
+    UP = "up"
+    DOWN = "down"
+
+
+@enum_tools.documentation.document_enum
+class Packing(str, Enum):
+    """
+    The packing method to be used
+    """
+
+    BIN = "bin"
+    STACK = "stack"
+
+
+class HierarchicalLayoutOptions(BaseModel, extra="forbid"):
+    """
+    The options for the hierarchical layout.
+    """
+
+    direction: Optional[Direction] = None
+    packaging: Optional[Packing] = None
+
+
+class ForceDirectedLayoutOptions(BaseModel, extra="forbid"):
+    """
+    The options for the force-directed layout.
+    """
+
+    gravity: Optional[float] = None
+    simulationStopVelocity: Optional[float] = None
+
+
+LayoutOptions = Union[HierarchicalLayoutOptions, ForceDirectedLayoutOptions]
+
+
+def construct_layout_options(layout: Layout, options: dict[str, Any]) -> Optional[LayoutOptions]:
+    if not options:
+        return None
+
+    if layout == Layout.FORCE_DIRECTED:
+        try:
+            return ForceDirectedLayoutOptions(**options)
+        except ValidationError as e:
+            _parse_validation_error(e, ForceDirectedLayoutOptions)
+    elif layout == Layout.HIERARCHICAL:
+        try:
+            return HierarchicalLayoutOptions(**options)
+        except ValidationError as e:
+            _parse_validation_error(e, ForceDirectedLayoutOptions)
+
+    raise ValueError(
+        f"Layout options only supported for layouts `{Layout.FORCE_DIRECTED}` and `{Layout.HIERARCHICAL}`, but was `{layout}`"
+    )
+
+
 @enum_tools.documentation.document_enum
 class Renderer(str, Enum):
     """
@@ -72,6 +135,9 @@ class RenderOptions(BaseModel, extra="allow"):
     """
 
     layout: Optional[Layout] = None
+    layout_options: Optional[Union[HierarchicalLayoutOptions, ForceDirectedLayoutOptions]] = Field(
+        None, serialization_alias="layoutOptions"
+    )
     renderer: Optional[Renderer] = None
 
     pan_X: Optional[float] = Field(None, serialization_alias="panX")
@@ -84,5 +150,34 @@ class RenderOptions(BaseModel, extra="allow"):
     min_zoom: Optional[float] = Field(None, serialization_alias="minZoom", description="The minimum zoom level allowed")
     allow_dynamic_min_zoom: Optional[bool] = Field(None, serialization_alias="allowDynamicMinZoom")
 
+    @model_validator(mode="after")
+    def check_layout_options_match(self) -> RenderOptions:
+        if self.layout_options is None:
+            return self
+
+        if self.layout == Layout.HIERARCHICAL and not isinstance(self.layout_options, HierarchicalLayoutOptions):
+            raise ValueError("layout_options must be of type HierarchicalLayoutOptions for hierarchical layout")
+        if self.layout == Layout.FORCE_DIRECTED and not isinstance(self.layout_options, ForceDirectedLayoutOptions):
+            raise ValueError("layout_options must be of type ForceDirectedLayoutOptions for force-directed layout")
+        return self
+
     def to_dict(self) -> dict[str, Any]:
         return self.model_dump(exclude_none=True, by_alias=True)
+
+
+def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) -> None:
+    for err in e.errors():
+        loc = err["loc"][0]
+        if err["type"] == "missing":
+            raise ValueError(
+                f"Mandatory `{entity_type.__name__}` parameter '{loc}' is missing. Expected one of {entity_type.model_fields[loc].validation_alias.choices} to be present"  # type: ignore
+            )
+        elif err["type"] == "extra_forbidden":
+            raise ValueError(
+                f"Unexpected `{entity_type.__name__}` parameter '{loc}' with provided input '{err['input']}'. "
+                f"Allowed parameters are: {', '.join(entity_type.model_fields.keys())}"
+            )
+        else:
+            raise ValueError(
+                f"Error for `{entity_type.__name__}` parameter '{loc}' with provided input '{err['input']}'. Reason: {err['msg']}"
+            )

python-wrapper/src/neo4j_viz/visualization_graph.py

@@ -2,7 +2,7 @@
 
 import warnings
 from collections.abc import Iterable
-from typing import Any, Callable, Hashable, Optional
+from typing import Any, Callable, Hashable, Optional, Union
 
 from IPython.display import HTML
 from pydantic_extra_types.color import Color, ColorType
@@ -11,7 +11,13 @@
 from .node import Node, NodeIdType
 from .node_size import RealNumber, verify_radii
 from .nvl import NVL
-from .options import Layout, Renderer, RenderOptions
+from .options import (
+    Layout,
+    LayoutOptions,
+    Renderer,
+    RenderOptions,
+    construct_layout_options,
+)
 from .relationship import Relationship
 
 
@@ -42,6 +48,7 @@ def __init__(self, nodes: list[Node], relationships: list[Relationship]) -> None
     def render(
         self,
         layout: Optional[Layout] = None,
+        layout_options: Union[dict[str, Any], LayoutOptions, None] = None,
         renderer: Renderer = Renderer.CANVAS,
         width: str = "100%",
         height: str = "600px",
@@ -60,6 +67,8 @@ def render(
         ----------
         layout:
             The `Layout` to use.
+        layout_options:
+            The `LayoutOptions` to use.
         renderer:
             The `Renderer` to use.
         width:
@@ -92,8 +101,19 @@ def render(
 
         Renderer.check(renderer, num_nodes)
 
+        if not layout:
+            layout = Layout.FORCE_DIRECTED
+        if not layout_options:
+            layout_options = {}
+
+        if isinstance(layout_options, dict):
+            layout_options_typed = construct_layout_options(layout, layout_options)
+        else:
+            layout_options_typed = layout_options
+
         render_options = RenderOptions(
             layout=layout,
+            layout_options=layout_options_typed,
             renderer=renderer,
             pan_X=pan_position[0] if pan_position is not None else None,
             pan_Y=pan_position[1] if pan_position is not None else None,

python-wrapper/tests/test_options.py

@@ -1,7 +1,24 @@
-from neo4j_viz.options import Layout, RenderOptions
+import pytest
+
+from neo4j_viz.options import Direction, HierarchicalLayoutOptions, Layout, RenderOptions
 
 
 def test_render_options() -> None:
-    options = RenderOptions(layout=Layout.FORCE_DIRECTED)
+    options = RenderOptions(layout=Layout.HIERARCHICAL)
+
+    assert options.to_dict() == {"layout": "hierarchical"}
+
+
+def test_render_options_with_layout_options() -> None:
+    options = RenderOptions(
+        layout=Layout.HIERARCHICAL, layout_options=HierarchicalLayoutOptions(direction=Direction.LEFT)
+    )
+
+    assert options.to_dict() == {"layout": "hierarchical", "layoutOptions": {"direction": "left"}}
+
 
-    assert options.to_dict() == {"layout": "forcedirected"}
+def test_layout_options_match() -> None:
+    with pytest.raises(
+        ValueError, match="layout_options must be of type ForceDirectedLayoutOptions for force-directed layout"
+    ):
+        RenderOptions(layout=Layout.FORCE_DIRECTED, layout_options=HierarchicalLayoutOptions(direction=Direction.LEFT))

python-wrapper/tests/test_render.py

@@ -14,6 +14,8 @@
     "force layout": {"layout": Layout.FORCE_DIRECTED},
     "grid layout": {"layout": Layout.GRID},
     "coordinate layout": {"layout": Layout.COORDINATE},
+    "hierarchical layout + options": {"layout": Layout.HIERARCHICAL, "layout_options": {"direction": "left"}},
+    "with layout options": {"layout_options": {"gravity": 0.1}},
 }
 
 
@@ -111,3 +113,23 @@ def test_render_non_json_serializable() -> None:
     VG = VisualizationGraph(nodes=[node], relationships=[])
     # Should not raise an error
     VG.render()
+
+
+def test_render_with_wrong_layout_options() -> None:
+    nodes = [
+        Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:0", caption="Person", x=1, y=10),
+    ]
+
+    VG = VisualizationGraph(nodes=nodes, relationships=[])
+
+    with pytest.raises(
+        ValueError,
+        match="Unexpected `ForceDirectedLayoutOptions` parameter 'direction' with provided input 'left'",
+    ):
+        VG.render(layout_options={"direction": "left"})
+
+    with pytest.raises(
+        ValueError,
+        match="Unexpected `ForceDirectedLayoutOptions` parameter 'direction' with provided input 'left'",
+    ):
+        VG.render(layout=Layout.FORCE_DIRECTED, layout_options={"direction": "left"})