neo4j-viz 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

neo4j_viz/gds.py

@@ -8,16 +8,26 @@ from uuid import uuid4
 import pandas as pd
 from graphdatascience import Graph, GraphDataScience
 
+from neo4j_viz.colors import NEO4J_COLORS_DISCRETE, ColorSpace
+
 from .pandas import _from_dfs
 from .visualization_graph import VisualizationGraph
 
 
 def _fetch_node_dfs(
-    gds: GraphDataScience, G: Graph, node_properties_by_label: dict[str, list[str]], node_labels: list[str]
+    gds: GraphDataScience,
+    G: Graph,
+    node_properties_by_label: dict[str, list[str]],
+    node_labels: list[str],
+    additional_db_node_properties: list[str],
 ) -> dict[str, pd.DataFrame]:
     return {
         lbl: gds.graph.nodeProperties.stream(
-            G, node_properties=node_properties_by_label[lbl], node_labels=[lbl], separate_property_columns=True
+            G,
+            node_properties=node_properties_by_label[lbl],
+            node_labels=[lbl],
+            separate_property_columns=True,
+            db_node_properties=additional_db_node_properties,
         )
         for lbl in node_labels
     }
@@ -47,17 +57,20 @@ def _fetch_rel_dfs(gds: GraphDataScience, G: Graph) -> list[pd.DataFrame]:
 def from_gds(
     gds: GraphDataScience,
     G: Graph,
-    size_property: Optional[str] = None,
-    additional_node_properties: Optional[list[str]] = None,
-    node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+    node_properties: Optional[list[str]] = None,
+    db_node_properties: Optional[list[str]] = None,
     max_node_count: int = 10_000,
 ) -> VisualizationGraph:
     """
     Create a VisualizationGraph from a GraphDataScience object and a Graph object.
 
-    All `additional_node_properties` will be included in the visualization graph.
-    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.
+    By default:
+
+    * the caption of a node will be based on its `labels`.
+    * the caption of a relationship will be based on its `relationshipType`.
+    * the color of nodes will be set based on their label, unless there are more than 12 unique labels.
+
+    All `node_properties` and `db_node_properties` will be included in the visualization graph under the `properties` field.
     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.
 
@@ -67,45 +80,36 @@ def from_gds(
         GraphDataScience object.
     G : Graph
         Graph object.
-    size_property : str, optional
-        Property to use for node size, by default None.
-    additional_node_properties : list[str], optional
+    node_properties : list[str], optional
         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.
+        properties from the Graph will be fetched.
+    db_node_properties : list[str], optional
+        Additional node properties to fetch from the database, by default None. Only works if the graph was projected from the database.
     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.
     """
+    if db_node_properties is None:
+        db_node_properties = []
+
     node_properties_from_gds = G.node_properties()
     assert isinstance(node_properties_from_gds, pd.Series)
     actual_node_properties: dict[str, list[str]] = cast(dict[str, list[str]], node_properties_from_gds.to_dict())
     all_actual_node_properties = list(chain.from_iterable(actual_node_properties.values()))
 
-    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()}'")
-
     node_properties_by_label_sets: dict[str, set[str]] = dict()
-    if additional_node_properties is None:
+    if node_properties is None:
         node_properties_by_label_sets = {k: set(v) for k, v in actual_node_properties.items()}
     else:
-        for prop in additional_node_properties:
+        for prop in node_properties:
             if prop not in all_actual_node_properties:
                 raise ValueError(f"There is no node property '{prop}' in graph '{G.name()}'")
 
         for label, props in actual_node_properties.items():
             node_properties_by_label_sets[label] = {
-                prop for prop in actual_node_properties[label] if prop in additional_node_properties
+                prop for prop in actual_node_properties[label] if prop in node_properties
             }
 
-    if size_property is not None:
-        # For some reason mypy are unable to understand that this is dict[str, set[str]]
-        for label, props in node_properties_by_label_sets.items():  # type: ignore
-            props.add(size_property)  # type: ignore
-
     node_properties_by_label = {k: list(v) for k, v in node_properties_by_label_sets.items()}
 
     node_count = G.node_count()
@@ -129,7 +133,9 @@ def from_gds(
             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())
+        node_dfs = _fetch_node_dfs(
+            gds, G_fetched, node_properties_by_label, G_fetched.node_labels(), db_node_properties
+        )
         if property_name is not None:
             for df in node_dfs.values():
                 df.drop(columns=[property_name], inplace=True)
@@ -145,14 +151,7 @@ def from_gds(
         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 all_actual_node_properties and size_property != "size":
-            node_props_df.rename(columns={"size": "__size"}, inplace=True)
-        if additional_node_properties is not None and size_property not in additional_node_properties:
-            node_props_df.rename(columns={size_property: "size"}, inplace=True)
-        else:
-            node_props_df["size"] = node_props_df[size_property]
+    node_props_df = pd.concat(node_dfs.values(), ignore_index=True, axis=0).drop_duplicates(subset=["nodeId"])
 
     for lbl, df in node_dfs.items():
         if "labels" in all_actual_node_properties:
@@ -164,22 +163,22 @@ def from_gds(
 
     node_df = node_props_df.merge(node_labels_df, on="nodeId")
 
-    if "caption" not in all_actual_node_properties:
-        node_df["caption"] = node_df["labels"].astype(str)
+    try:
+        VG = _from_dfs(node_df, rel_dfs, dropna=True)
 
-    for rel_df in rel_dfs:
-        if "caption" not in rel_df.columns:
-            rel_df["caption"] = rel_df["relationshipType"]
+        for node in VG.nodes:
+            node.caption = ":".join([label for label in node.properties["labels"]])
+        for rel in VG.relationships:
+            rel.caption = rel.properties.get("relationshipType")
 
-    try:
-        return _from_dfs(
-            node_df, rel_dfs, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"}, dropna=True
-        )
+        number_of_colors = node_df["labels"].drop_duplicates().count()
+        if number_of_colors <= len(NEO4J_COLORS_DISCRETE):
+            VG.color_nodes(property="labels", color_space=ColorSpace.DISCRETE)
+
+        return VG
     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

@@ -5,6 +5,7 @@ from typing import Any, Optional
 from pydantic import BaseModel, ValidationError
 
 from neo4j_viz import Node, Relationship, VisualizationGraph
+from neo4j_viz.colors import NEO4J_COLORS_DISCRETE, ColorSpace
 
 
 def _parse_value(value_str: str) -> Any:
@@ -91,10 +92,7 @@ def _parse_value(value_str: str) -> Any:
     return value_str.strip("'\"")
 
 
-def _parse_prop_str(
-    query: str, prop_str: str, prop_start: int, top_level_keys: set[str]
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    top_level: dict[str, Any] = {}
+def _parse_prop_str(query: str, prop_str: str, prop_start: int) -> dict[str, Any]:
     props: dict[str, Any] = {}
     depth = 0
     in_string = None
@@ -115,10 +113,7 @@ def _parse_prop_str(
                 k, v = pair.split(":", 1)
                 k = k.strip().strip("'\"")
 
-                if k in top_level_keys:
-                    top_level[k] = _parse_value(v)
-                else:
-                    props[k] = _parse_value(v)
+                props[k] = _parse_value(v)
 
                 start_idx = i + 1
         else:
@@ -133,17 +128,12 @@ def _parse_prop_str(
         k, v = pair.split(":", 1)
         k = k.strip().strip("'\"")
 
-        if k in top_level_keys:
-            top_level[k] = _parse_value(v)
-        else:
-            props[k] = _parse_value(v)
+        props[k] = _parse_value(v)
 
-    return top_level, props
+    return props
 
 
-def _parse_labels_and_props(
-    query: str, s: str, top_level_keys: set[str]
-) -> tuple[Optional[str], dict[str, Any], dict[str, Any]]:
+def _parse_labels_and_props(query: str, s: str) -> tuple[Optional[str], dict[str, Any]]:
     prop_match = re.search(r"\{(.*)\}", s)
     prop_str = ""
     if prop_match:
@@ -155,9 +145,8 @@ def _parse_labels_and_props(
     final_alias = raw_alias if raw_alias else None
 
     if prop_str:
-        top_level, props = _parse_prop_str(query, prop_str, prop_start, top_level_keys)
+        props = _parse_prop_str(query, prop_str, prop_start)
     else:
-        top_level = {}
         props = {}
 
     label_list = [lbl.strip() for lbl in alias_labels[1:]]
@@ -165,7 +154,7 @@ def _parse_labels_and_props(
         props["__labels"] = props["labels"]
     props["labels"] = sorted(label_list)
 
-    return final_alias, top_level, props
+    return final_alias, props
 
 
 def _get_snippet(q: str, idx: int, context: int = 15) -> str:
@@ -175,21 +164,20 @@ def _get_snippet(q: str, idx: int, context: int = 15) -> str:
     return q[start:end].replace("\n", " ")
 
 
-def from_gql_create(
-    query: str,
-    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),
-) -> VisualizationGraph:
+def from_gql_create(query: str) -> VisualizationGraph:
     """
     Parse a GQL CREATE query and return a VisualizationGraph object representing the graph it creates.
 
     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
-    top level fields of the respective objects. Otherwise, they will be included in the `properties` dictionary.
+    All properties of nodes and relationships will be included in the `properties` dictionary of the respective objects.
     Additionally, a "labels" property will be added for nodes and a "type" property for relationships.
 
+    By default:
+
+    * the caption of a node will be based on its `labels`.
+    * the caption of a relationship will be based on its `type`.
+    * the color of nodes will be set based on their label, unless there are more than 12 unique labels.
+
     Please note that this function is not a full GQL parser, it only handles CREATE queries that do not contain
     other clauses like MATCH, WHERE, RETURN, etc, or any Cypher function calls.
     It also does not handle all possible GQL syntax, but it should work for most common cases.
@@ -199,15 +187,6 @@ def from_gql_create(
     ----------
     query : str
         The GQL CREATE query to parse
-    size_property : str, optional
-        Property to use for node size, by default None.
-    node_caption : str, optional
-        Property to use as the node caption, by default the node labels will be used.
-    relationship_caption : str, optional
-        Property to use as the relationship caption, by default the relationship type will be used.
-    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.
     """
 
     query = query.strip()
@@ -251,19 +230,9 @@ def from_gql_create(
     node_pattern = re.compile(r"^\(([^)]*)\)$")
     rel_pattern = re.compile(r"^\(([^)]*)\)-\s*\[\s*:(\w+)\s*(\{[^}]*\})?\s*\]->\(([^)]*)\)$")
 
-    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']}"
             )
@@ -277,14 +246,14 @@ def from_gql_create(
         node_m = node_pattern.match(part)
         if node_m:
             alias_labels_props = node_m.group(1).strip()
-            alias, top_level, props = _parse_labels_and_props(query, alias_labels_props, node_top_level_keys)
+            alias, props = _parse_labels_and_props(query, alias_labels_props)
             if not alias:
                 alias = f"_anon_{anonymous_count}"
                 anonymous_count += 1
             if alias not in alias_to_id:
                 alias_to_id[alias] = str(uuid.uuid4())
             try:
-                nodes.append(Node(id=alias_to_id[alias], **top_level, properties=props))
+                nodes.append(Node(id=alias_to_id[alias], properties=props))
             except ValidationError as e:
                 _parse_validation_error(e, Node)
 
@@ -296,14 +265,14 @@ def from_gql_create(
             right_node = rel_m.group(4).strip()
 
             # Parse left node pattern
-            left_alias, left_top_level, left_props = _parse_labels_and_props(query, left_node, node_top_level_keys)
+            left_alias, left_props = _parse_labels_and_props(query, left_node)
             if not left_alias:
                 left_alias = f"_anon_{anonymous_count}"
                 anonymous_count += 1
                 if left_alias not in alias_to_id:
                     alias_to_id[left_alias] = str(uuid.uuid4())
                 try:
-                    nodes.append(Node(id=alias_to_id[left_alias], **left_top_level, properties=left_props))
+                    nodes.append(Node(id=alias_to_id[left_alias], properties=left_props))
                 except ValidationError as e:
                     _parse_validation_error(e, Node)
             elif left_alias not in alias_to_id:
@@ -311,14 +280,14 @@ def from_gql_create(
                 raise ValueError(f"Relationship references unknown node alias: '{left_alias}' near: `{snippet}`.")
 
             # Parse right node pattern
-            right_alias, right_top_level, right_props = _parse_labels_and_props(query, right_node, node_top_level_keys)
+            right_alias, right_props = _parse_labels_and_props(query, right_node)
             if not right_alias:
                 right_alias = f"_anon_{anonymous_count}"
                 anonymous_count += 1
                 if right_alias not in alias_to_id:
                     alias_to_id[right_alias] = str(uuid.uuid4())
                 try:
-                    nodes.append(Node(id=alias_to_id[right_alias], **right_top_level, properties=right_props))
+                    nodes.append(Node(id=alias_to_id[right_alias], properties=right_props))
                 except ValidationError as e:
                     _parse_validation_error(e, Node)
             elif right_alias not in alias_to_id:
@@ -331,9 +300,8 @@ def from_gql_create(
             if rel_props_str:
                 inner_str = rel_props_str.strip("{}").strip()
                 prop_start = query.index(inner_str, query.index(inner_str))
-                top_level, props = _parse_prop_str(query, inner_str, prop_start, rel_top_level_keys)
+                props = _parse_prop_str(query, inner_str, prop_start)
             else:
-                top_level = {}
                 props = {}
             if "type" in props:
                 props["__type"] = props["type"]
@@ -345,7 +313,6 @@ def from_gql_create(
                         id=rel_id,
                         source=alias_to_id[left_alias],
                         target=alias_to_id[right_alias],
-                        **top_level,
                         properties=props,
                     )
                 )
@@ -357,29 +324,15 @@ def from_gql_create(
         snippet = part[:30]
         raise ValueError(f"Invalid element in CREATE near: `{snippet}`.")
 
-    if size_property is not None:
-        for node in nodes:
-            node.size = node.properties.get(size_property)
-    if node_caption is not None:
-        for node in nodes:
-            if node_caption == "labels":
-                if len(node.properties["labels"]) > 0:
-                    node.caption = ":".join([label for label in node.properties["labels"]])
-            else:
-                node.caption = str(node.properties.get(node_caption))
-    if relationship_caption is not None:
-        for rel in relationships:
-            if relationship_caption == "type":
-                rel.caption = rel.properties["type"]
-            else:
-                rel.caption = str(rel.properties.get(relationship_caption))
-
     VG = VisualizationGraph(nodes=nodes, relationships=relationships)
-    if (node_radius_min_max is not None) and (size_property is not None):
-        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")
+
+    for node in VG.nodes:
+        node.caption = ":".join([label for label in node.properties["labels"]])
+    for rel in VG.relationships:
+        rel.caption = rel.properties.get("type")
+
+    number_of_colors = len({str(n.properties.get("labels")) for n in VG.nodes})
+    if number_of_colors <= len(NEO4J_COLORS_DISCRETE):
+        VG.color_nodes(property="labels", color_space=ColorSpace.DISCRETE)
 
     return VG

neo4j_viz/neo4j.py

@@ -7,6 +7,7 @@ import neo4j.graph
 from neo4j import Driver, Result, RoutingControl
 from pydantic import BaseModel, ValidationError
 
+from neo4j_viz.colors import NEO4J_COLORS_DISCRETE, ColorSpace
 from neo4j_viz.node import Node
 from neo4j_viz.relationship import Relationship
 from neo4j_viz.visualization_graph import VisualizationGraph
@@ -22,18 +23,18 @@ def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) ->
 
 def from_neo4j(
     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`, 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
-    top level fields of the respective objects. Otherwise, they will be included in the `properties` dictionary.
+    By default:
+
+    * the caption of a node will be based on its `labels`.
+    * the caption of a relationship will be based on its `type`.
+    * the color of nodes will be set based on their label, unless there are more than 12 unique labels.
+
+    All node and relationship properties will be included in the visualization graph under the `properties` field.
     Additionally, a "labels" property will be added for nodes and a "type" property for relationships.
 
     Parameters
@@ -41,15 +42,6 @@ def from_neo4j(
     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
-        Property to use as the node caption, by default the node labels will be used.
-    relationship_caption : str, optional
-        Property to use as the relationship caption, by default the relationship type will be used.
-    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.
@@ -77,117 +69,62 @@ def from_neo4j(
     else:
         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()
-
-    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)
+    nodes = [_map_node(node) for node in graph.nodes]
 
     relationships = []
-    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)
+
+    for rel in graph.relationships:
+        mapped_rel = _map_relationship(rel)
+        if mapped_rel:
+            relationships.append(mapped_rel)
 
     VG = VisualizationGraph(nodes, 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)
+    for node in VG.nodes:
+        node.caption = ":".join(node.properties["labels"])
+    for r in VG.relationships:
+        r.caption = r.properties["type"]
+
+    number_of_colors = len({n.caption for n in VG.nodes})
+    if number_of_colors <= len(NEO4J_COLORS_DISCRETE):
+        VG.color_nodes(field="caption", color_space=ColorSpace.DISCRETE, colors=NEO4J_COLORS_DISCRETE)
 
     return VG
 
 
 def _map_node(
     node: neo4j.graph.Node,
-    all_node_field_aliases: set[str],
-    size_property: Optional[str],
-    caption_property: Optional[str],
 ) -> Node:
-    top_level_fields = {"id": node.element_id}
-
-    if size_property:
-        top_level_fields["size"] = node.get(size_property)
-
     labels = sorted([label for label in node.labels])
-    if caption_property:
-        if caption_property == "labels":
-            if len(labels) > 0:
-                top_level_fields["caption"] = ":".join([label for label in labels])
-        else:
-            top_level_fields["caption"] = str(node.get(caption_property))
-
-    properties = {}
-    for prop, value in node.items():
-        if prop not in all_node_field_aliases:
-            properties[prop] = value
-            continue
 
-        if prop in top_level_fields:
-            properties[prop] = value
-            continue
-
-        top_level_fields[prop] = value
+    properties = {prop: value for prop, value in node.items()}
 
     if "labels" in properties:
         properties["__labels"] = properties["labels"]
     properties["labels"] = labels
 
     try:
-        viz_node = Node(**top_level_fields, properties=properties)
+        viz_node = Node(id=node.element_id, properties=properties)
     except ValidationError as e:
         _parse_validation_error(e, Node)
 
     return viz_node
 
 
-def _map_relationship(
-    rel: neo4j.graph.Relationship, all_rel_field_aliases: set[str], caption_property: Optional[str]
-) -> Optional[Relationship]:
+def _map_relationship(rel: neo4j.graph.Relationship) -> Optional[Relationship]:
     if rel.start_node is None or rel.end_node is None:
         return None
 
-    top_level_fields = {"id": rel.element_id, "source": rel.start_node.element_id, "target": rel.end_node.element_id}
-
-    if caption_property:
-        if caption_property == "type":
-            top_level_fields["caption"] = rel.type
-        else:
-            top_level_fields["caption"] = str(rel.get(caption_property))
-
-    properties = {}
-    for prop, value in rel.items():
-        if prop not in all_rel_field_aliases:
-            properties[prop] = value
-            continue
-
-        if prop in top_level_fields:
-            properties[prop] = value
-            continue
-
-        top_level_fields[prop] = value
+    properties = {prop: value for prop, value in rel.items()}
 
     if "type" in properties:
         properties["__type"] = properties["type"]
     properties["type"] = rel.type
 
     try:
-        viz_rel = Relationship(**top_level_fields, properties=properties)
+        viz_rel = Relationship(
+            id=rel.element_id, source=rel.start_node.element_id, target=rel.end_node.element_id, properties=properties
+        )
     except ValidationError as e:
         _parse_validation_error(e, Relationship)
 

neo4j_viz/node.py

@@ -30,6 +30,7 @@ class Node(
         validation_alias=create_aliases,
         serialization_alias=lambda field_name: to_camel(field_name),
     ),
+    validate_assignment=True,
 ):
     """
     A node in a graph to visualize.
@@ -90,10 +91,8 @@ class Node(
         return self.model_dump(exclude_none=True, by_alias=True)
 
     @staticmethod
-    def all_validation_aliases(exempted_fields: Optional[list[str]] = None) -> set[str]:
-        if exempted_fields is None:
-            exempted_fields = []
-
-        by_field = [v.validation_alias.choices for k, v in Node.model_fields.items() if k not in exempted_fields]  # type: ignore
+    def basic_fields_validation_aliases() -> set[str]:
+        mandatory_fields = ["id"]
+        by_field = [v.validation_alias.choices for k, v in Node.model_fields.items() if k in mandatory_fields]  # type: ignore
 
         return {str(alias) for aliases in by_field for alias in aliases}

neo4j_viz/options.py

@@ -22,6 +22,9 @@ class CaptionAlignment(str, Enum):
 @enum_tools.documentation.document_enum
 class Layout(str, Enum):
     FORCE_DIRECTED = "forcedirected"
+    """
+    The force-directed layout uses a physics simulation to position the nodes.
+    """
     HIERARCHICAL = "hierarchical"
     """
     The nodes are then arranged by the directionality of their relationships