neo4j-viz 0.3.1__py3-none-any.whl → 0.4.1__py3-none-any.whl

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",
+]

neo4j_viz/gds.py

@@ -1,27 +1,44 @@
 from __future__ import annotations
 
+import warnings
 from itertools import chain
 from typing import Optional
+from uuid import uuid4
 
 import pandas as pd
 from graphdatascience import Graph, GraphDataScience
+from pandas import Series
 
 from .pandas import _from_dfs
 from .visualization_graph import VisualizationGraph
 
 
-def _node_dfs(
-    gds: GraphDataScience, G: Graph, node_properties: list[str], node_labels: list[str]
+def _fetch_node_dfs(
+    gds: GraphDataScience, G: Graph, node_properties_by_label: dict[str, list[str]], node_labels: list[str]
 ) -> dict[str, pd.DataFrame]:
     return {
         lbl: gds.graph.nodeProperties.stream(
-            G, node_properties=node_properties, node_labels=[lbl], separate_property_columns=True
+            G, node_properties=node_properties_by_label[lbl], node_labels=[lbl], separate_property_columns=True
         )
         for lbl in node_labels
     }
 
 
-def _rel_df(gds: GraphDataScience, G: Graph) -> pd.DataFrame:
+def _fetch_rel_df(gds: GraphDataScience, G: Graph) -> pd.DataFrame:
+    relationship_properties = G.relationship_properties()
+    assert isinstance(relationship_properties, Series)
+
+    relationship_properties_per_type = relationship_properties.tolist()
+    property_set: set[str] = set()
+    for props in relationship_properties_per_type:
+        if props:
+            property_set.update(props)
+
+    if len(property_set) > 0:
+        return gds.graph.relationshipProperties.stream(
+            G, relationship_properties=list(property_set), separate_property_columns=True
+        )
+
     return gds.graph.relationships.stream(G)
 
 
@@ -31,6 +48,7 @@ def from_gds(
     size_property: Optional[str] = None,
     additional_node_properties: Optional[list[str]] = None,
     node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+    max_node_count: int = 10_000,
 ) -> VisualizationGraph:
     """
     Create a VisualizationGraph from a GraphDataScience object and a Graph object.
@@ -39,6 +57,7 @@ def from_gds(
     If the properties are named as the fields of the `Node` class, they will be included as top level fields of the
     created `Node` objects. Otherwise, they will be included in the `properties` dictionary.
     Additionally, a new "labels" node property will be added, containing the node labels of the node.
+    Similarly for relationships, a new "relationshipType" property will be added.
 
     Parameters
     ----------
@@ -49,52 +68,111 @@ def from_gds(
     size_property : str, optional
         Property to use for node size, by default None.
     additional_node_properties : list[str], optional
-        Additional properties to include in the visualization node, by default None. They can be used later for modifying the node appearance.
+        Additional properties to include in the visualization node, by default None which means that all node
+        properties will be fetched.
     node_radius_min_max : tuple[float, float], optional
         Minimum and maximum node radius, by default (3, 60).
         To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
+    max_node_count : int, optional
+        The maximum number of nodes to fetch from the graph. The graph will be sampled using random walk with restarts
+        if its node count exceeds this number.
     """
     node_properties_from_gds = G.node_properties()
     assert isinstance(node_properties_from_gds, pd.Series)
-    actual_node_properties = list(chain.from_iterable(node_properties_from_gds.to_dict().values()))
+    actual_node_properties = node_properties_from_gds.to_dict()
+    all_actual_node_properties = list(chain.from_iterable(actual_node_properties.values()))
 
-    if size_property is not None and size_property not in actual_node_properties:
-        raise ValueError(f"There is no node property '{size_property}' in graph '{G.name()}'")
+    if size_property is not None:
+        if size_property not in all_actual_node_properties:
+            raise ValueError(f"There is no node property '{size_property}' in graph '{G.name()}'")
 
-    if additional_node_properties is not None:
+    if additional_node_properties is None:
+        node_properties_by_label = {k: set(v) for k, v in actual_node_properties.items()}
+    else:
         for prop in additional_node_properties:
-            if prop not in actual_node_properties:
+            if prop not in all_actual_node_properties:
                 raise ValueError(f"There is no node property '{prop}' in graph '{G.name()}'")
 
-    node_properties = set()
-    if additional_node_properties is not None:
-        node_properties.update(additional_node_properties)
+        node_properties_by_label = {}
+        for label, props in actual_node_properties.items():
+            node_properties_by_label[label] = {
+                prop for prop in actual_node_properties[label] if prop in additional_node_properties
+            }
 
     if size_property is not None:
-        node_properties.add(size_property)
+        for label, props in node_properties_by_label.items():
+            props.add(size_property)
+
+    node_properties_by_label = {k: list(v) for k, v in node_properties_by_label.items()}
+
+    node_count = G.node_count()
+    if node_count > max_node_count:
+        warnings.warn(
+            f"The '{G.name()}' projection's node count ({G.node_count()}) exceeds `max_node_count` ({max_node_count}), so subsampling will be applied. Increase `max_node_count` if needed"
+        )
+        sampling_ratio = float(max_node_count) / node_count
+        sample_name = f"neo4j-viz_sample_{uuid4()}"
+        G_fetched, _ = gds.graph.sample.rwr(sample_name, G, samplingRatio=sampling_ratio, nodeLabelStratification=True)
+    else:
+        G_fetched = G
+
+    property_name = None
+    try:
+        # Since GDS does not allow us to only fetch node IDs, we add the degree property
+        # as a temporary property to ensure that we have at least one property for each label to fetch
+        if sum([len(props) == 0 for props in node_properties_by_label.values()]) > 0:
+            property_name = f"neo4j-viz_property_{uuid4()}"
+            gds.degree.mutate(G_fetched, mutateProperty=property_name)
+            for props in node_properties_by_label.values():
+                props.append(property_name)
+
+        node_dfs = _fetch_node_dfs(gds, G_fetched, node_properties_by_label, G_fetched.node_labels())
+        if property_name is not None:
+            for df in node_dfs.values():
+                df.drop(columns=[property_name], inplace=True)
+
+        rel_df = _fetch_rel_df(gds, G_fetched)
+    finally:
+        if G_fetched.name() != G.name():
+            G_fetched.drop()
+        elif property_name is not None:
+            gds.graph.nodeProperties.drop(G_fetched, node_properties=[property_name])
 
-    node_properties = list(node_properties)
-    node_dfs = _node_dfs(gds, G, node_properties, G.node_labels())
     for df in node_dfs.values():
-        df.rename(columns={"nodeId": "id"}, inplace=True)
+        if property_name is not None and property_name in df.columns:
+            df.drop(columns=[property_name], inplace=True)
 
     node_props_df = pd.concat(node_dfs.values(), ignore_index=True, axis=0).drop_duplicates()
     if size_property is not None:
-        if "size" in actual_node_properties and size_property != "size":
+        if "size" in all_actual_node_properties and size_property != "size":
             node_props_df.rename(columns={"size": "__size"}, inplace=True)
         node_props_df.rename(columns={size_property: "size"}, inplace=True)
 
     for lbl, df in node_dfs.items():
-        if "labels" in actual_node_properties:
+        if "labels" in all_actual_node_properties:
             df.rename(columns={"labels": "__labels"}, inplace=True)
         df["labels"] = lbl
 
-    node_lbls_df = pd.concat([df[["id", "labels"]] for df in node_dfs.values()], ignore_index=True, axis=0)
-    node_lbls_df = node_lbls_df.groupby("id").agg({"labels": list})
+    node_labels_df = pd.concat([df[["nodeId", "labels"]] for df in node_dfs.values()], ignore_index=True, axis=0)
+    node_labels_df = node_labels_df.groupby("nodeId").agg({"labels": list})
 
-    node_df = node_props_df.merge(node_lbls_df, on="id")
+    node_df = node_props_df.merge(node_labels_df, on="nodeId")
 
-    rel_df = _rel_df(gds, G)
-    rel_df.rename(columns={"sourceNodeId": "source", "targetNodeId": "target"}, inplace=True)
+    if "caption" not in all_actual_node_properties:
+        node_df["caption"] = node_df["labels"].astype(str)
 
-    return _from_dfs(node_df, rel_df, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"})
+    if "caption" not in rel_df.columns:
+        rel_df["caption"] = rel_df["relationshipType"]
+
+    try:
+        return _from_dfs(
+            node_df, rel_df, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"}, dropna=True
+        )
+    except ValueError as e:
+        err_msg = str(e)
+        if "column" in err_msg:
+            err_msg = err_msg.replace("column", "property")
+            if ("'size'" in err_msg) and (size_property is not None):
+                err_msg = err_msg.replace("'size'", f"'{size_property}'")
+            raise ValueError(err_msg)
+        raise e

neo4j_viz/gql_create.py

@@ -2,6 +2,8 @@ import re
 import uuid
 from typing import Any, Optional
 
+from pydantic import BaseModel, ValidationError
+
 from neo4j_viz import Node, Relationship, VisualizationGraph
 
 
@@ -252,6 +254,20 @@ def from_gql_create(
     node_top_level_keys = Node.all_validation_aliases(exempted_fields=["id"])
     rel_top_level_keys = Relationship.all_validation_aliases(exempted_fields=["id", "source", "target"])
 
+    def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) -> None:
+        for err in e.errors():
+            loc = err["loc"][0]
+            if (loc == "size") and size_property is not None:
+                loc = size_property
+            if loc == "caption":
+                if (entity_type == Node) and (node_caption is not None):
+                    loc = node_caption
+                elif (entity_type == Relationship) and (relationship_caption is not None):
+                    loc = relationship_caption
+            raise ValueError(
+                f"Error for {entity_type.__name__.lower()} property '{loc}' with provided input '{err['input']}'. Reason: {err['msg']}"
+            )
+
     nodes = []
     relationships = []
     alias_to_id = {}
@@ -267,7 +283,10 @@ def from_gql_create(
                 anonymous_count += 1
             if alias not in alias_to_id:
                 alias_to_id[alias] = str(uuid.uuid4())
-            nodes.append(Node(id=alias_to_id[alias], **top_level, properties=props))
+            try:
+                nodes.append(Node(id=alias_to_id[alias], **top_level, properties=props))
+            except ValidationError as e:
+                _parse_validation_error(e, Node)
 
             continue
 
@@ -283,7 +302,10 @@ def from_gql_create(
                 anonymous_count += 1
                 if left_alias not in alias_to_id:
                     alias_to_id[left_alias] = str(uuid.uuid4())
-                nodes.append(Node(id=alias_to_id[left_alias], **left_top_level, properties=left_props))
+                try:
+                    nodes.append(Node(id=alias_to_id[left_alias], **left_top_level, properties=left_props))
+                except ValidationError as e:
+                    _parse_validation_error(e, Node)
             elif left_alias not in alias_to_id:
                 snippet = _get_snippet(query, query.index(left_node))
                 raise ValueError(f"Relationship references unknown node alias: '{left_alias}' near: `{snippet}`.")
@@ -295,7 +317,10 @@ def from_gql_create(
                 anonymous_count += 1
                 if right_alias not in alias_to_id:
                     alias_to_id[right_alias] = str(uuid.uuid4())
-                nodes.append(Node(id=alias_to_id[right_alias], **right_top_level, properties=right_props))
+                try:
+                    nodes.append(Node(id=alias_to_id[right_alias], **right_top_level, properties=right_props))
+                except ValidationError as e:
+                    _parse_validation_error(e, Node)
             elif right_alias not in alias_to_id:
                 snippet = _get_snippet(query, query.index(right_node))
                 raise ValueError(f"Relationship references unknown node alias: '{right_alias}' near: `{snippet}`.")
@@ -313,15 +338,20 @@ def from_gql_create(
             if "type" in props:
                 props["__type"] = props["type"]
             props["type"] = rel_type
-            relationships.append(
-                Relationship(
-                    id=rel_id,
-                    source=alias_to_id[left_alias],
-                    target=alias_to_id[right_alias],
-                    **top_level,
-                    properties=props,
+
+            try:
+                relationships.append(
+                    Relationship(
+                        id=rel_id,
+                        source=alias_to_id[left_alias],
+                        target=alias_to_id[right_alias],
+                        **top_level,
+                        properties=props,
+                    )
                 )
-            )
+            except ValidationError as e:
+                _parse_validation_error(e, Relationship)
+
             continue
 
         snippet = part[:30]
@@ -346,6 +376,10 @@ def from_gql_create(
 
     VG = VisualizationGraph(nodes=nodes, relationships=relationships)
     if (node_radius_min_max is not None) and (size_property is not None):
-        VG.resize_nodes(node_radius_min_max=node_radius_min_max)
+        try:
+            VG.resize_nodes(node_radius_min_max=node_radius_min_max)
+        except TypeError:
+            loc = "size" if size_property is None else size_property
+            raise ValueError(f"Error for node property '{loc}'. Reason: must be a numerical value")
 
     return VG

neo4j_viz/neo4j.py

@@ -1,24 +1,35 @@
 from __future__ import annotations
 
+import warnings
 from typing import Optional, Union
 
 import neo4j.graph
-from neo4j import Result
+from neo4j import Driver, Result, RoutingControl
+from pydantic import BaseModel, ValidationError
 
 from neo4j_viz.node import Node
 from neo4j_viz.relationship import Relationship
 from neo4j_viz.visualization_graph import VisualizationGraph
 
 
+def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) -> None:
+    for err in e.errors():
+        loc = err["loc"][0]
+        raise ValueError(
+            f"Error for {entity_type.__name__.lower()} property '{loc}' with provided input '{err['input']}'. Reason: {err['msg']}"
+        )
+
+
 def from_neo4j(
-    result: Union[neo4j.graph.Graph, Result],
+    data: Union[neo4j.graph.Graph, Result, Driver],
     size_property: Optional[str] = None,
     node_caption: Optional[str] = "labels",
     relationship_caption: Optional[str] = "type",
     node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+    row_limit: int = 10_000,
 ) -> VisualizationGraph:
     """
-    Create a VisualizationGraph from a Neo4j Graph or Neo4j Result object.
+    Create a VisualizationGraph from a Neo4j `Graph`, Neo4j `Result` or Neo4j `Driver`.
 
     All node and relationship properties will be included in the visualization graph.
     If the properties are named as the fields of the `Node` or `Relationship` classes, they will be included as
@@ -27,8 +38,9 @@ def from_neo4j(
 
     Parameters
     ----------
-    result : Union[neo4j.graph.Graph, Result]
-        Query result either in shape of a Graph or result.
+    data : Union[neo4j.graph.Graph, neo4j.Result, neo4j.Driver]
+        Either a query result in the shape of a `neo4j.graph.Graph` or `neo4j.Result`, or a `neo4j.Driver` in
+        which case a simple default query will be executed internally to retrieve the graph data.
     size_property : str, optional
         Property to use for node size, by default None.
     node_caption : str, optional
@@ -38,26 +50,60 @@ def from_neo4j(
     node_radius_min_max : tuple[float, float], optional
         Minimum and maximum node radius, by default (3, 60).
         To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
+    row_limit : int, optional
+        Maximum number of rows to return from the query, by default 10_000.
+        This is only used if a `neo4j.Driver` is passed as `result` argument, otherwise the limit is ignored.
     """
 
-    if isinstance(result, Result):
-        graph = result.graph()
-    elif isinstance(result, neo4j.graph.Graph):
-        graph = result
+    if isinstance(data, Result):
+        graph = data.graph()
+    elif isinstance(data, neo4j.graph.Graph):
+        graph = data
+    elif isinstance(data, Driver):
+        rel_count = data.execute_query(
+            "MATCH ()-[r]->() RETURN count(r) as count",
+            routing_=RoutingControl.READ,
+            result_transformer_=Result.single,
+        ).get("count")  # type: ignore[union-attr]
+        if rel_count > row_limit:
+            warnings.warn(
+                f"Database relationship count ({rel_count}) exceeds `row_limit` ({row_limit}), so limiting will be applied. Increase the `row_limit` if needed"
+            )
+        graph = data.execute_query(
+            f"MATCH (n)-[r]->(m) RETURN n,r,m LIMIT {row_limit}",
+            routing_=RoutingControl.READ,
+            result_transformer_=Result.graph,
+        )
     else:
-        raise ValueError(f"Invalid input type `{type(result)}`. Expected `neo4j.Graph` or `neo4j.Result`")
+        raise ValueError(f"Invalid input type `{type(data)}`. Expected `neo4j.Graph`, `neo4j.Result` or `neo4j.Driver`")
 
     all_node_field_aliases = Node.all_validation_aliases()
     all_rel_field_aliases = Relationship.all_validation_aliases()
 
-    nodes = [
-        _map_node(node, all_node_field_aliases, size_property, caption_property=node_caption) for node in graph.nodes
-    ]
+    try:
+        nodes = [
+            _map_node(node, all_node_field_aliases, size_property, caption_property=node_caption)
+            for node in graph.nodes
+        ]
+    except ValueError as e:
+        err_msg = str(e)
+        if ("'size'" in err_msg) and (size_property is not None):
+            err_msg = err_msg.replace("'size'", f"'{size_property}'")
+        elif ("'caption'" in err_msg) and (node_caption is not None):
+            err_msg = err_msg.replace("'caption'", f"'{node_caption}'")
+        raise ValueError(err_msg)
+
     relationships = []
-    for rel in graph.relationships:
-        mapped_rel = _map_relationship(rel, all_rel_field_aliases, caption_property=relationship_caption)
-        if mapped_rel:
-            relationships.append(mapped_rel)
+    try:
+        for rel in graph.relationships:
+            mapped_rel = _map_relationship(rel, all_rel_field_aliases, caption_property=relationship_caption)
+            if mapped_rel:
+                relationships.append(mapped_rel)
+    except ValueError as e:
+        err_msg = str(e)
+        if ("'caption'" in err_msg) and (relationship_caption is not None):
+            err_msg = err_msg.replace("'caption'", f"'{relationship_caption}'")
+        raise ValueError(err_msg)
 
     VG = VisualizationGraph(nodes, relationships)
 
@@ -102,7 +148,12 @@ def _map_node(
         properties["__labels"] = properties["labels"]
     properties["labels"] = labels
 
-    return Node(**top_level_fields, properties=properties)
+    try:
+        viz_node = Node(**top_level_fields, properties=properties)
+    except ValidationError as e:
+        _parse_validation_error(e, Node)
+
+    return viz_node
 
 
 def _map_relationship(
@@ -135,4 +186,9 @@ def _map_relationship(
         properties["__type"] = properties["type"]
     properties["type"] = rel.type
 
-    return Relationship(**top_level_fields, properties=properties)
+    try:
+        viz_rel = Relationship(**top_level_fields, properties=properties)
+    except ValidationError as e:
+        _parse_validation_error(e, Relationship)
+
+    return viz_rel

neo4j_viz/options.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 
 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']}"
+            )