Merge pull request #53 from neo4j/streamlit-example

Add streamlit example

Author: adamnsch

examples/datasets/cora/cora_nodes.parquet.gzip

@@ -0,0 +1,65 @@
+import streamlit as st
+from IPython.display import HTML
+import streamlit.components.v1 as components
+from pandas import read_parquet
+import pathlib
+
+from neo4j_viz.pandas import from_dfs
+from neo4j_viz import VisualizationGraph
+
+# Path to this file
+script_path = pathlib.Path(__file__).resolve()
+script_dir_path = pathlib.Path(__file__).parent.resolve()
+
+
+@st.cache_data
+def create_visualization_graph() -> VisualizationGraph:
+    cora_nodes_path = f"{script_dir_path}/datasets/cora/cora_nodes.parquet.gzip"
+    cora_rels_path = f"{script_dir_path}/datasets/cora/cora_rels.parquet.gzip"
+
+    nodes_df = read_parquet(cora_nodes_path)
+    rels_df = read_parquet(cora_rels_path)
+
+    # Drop the features column since it's not needed for visualization
+    # Also numpy arrays are not supported by the visualization library
+    nodes_df.drop(columns="features", inplace=True)
+
+    VG = from_dfs(nodes_df, rels_df)
+    VG.color_nodes("subject")
+
+    return VG
+
+
+@st.cache_data
+def render_graph(
+    _VG: VisualizationGraph, height: int, initial_zoom: float = 0.1
+) -> HTML:
+    return VG.render(initial_zoom=initial_zoom, height=f"{height}px")
+
+
+VG = create_visualization_graph()
+
+st.title("Neo4j Viz Streamlit Example")
+st.text(
+    "This is an example of how to use Streamlit with the Graph "
+    "Visualization for Python library by Neo4j."
+)
+
+with st.sidebar:
+    height = st.slider("Height in pixels", 100, 2000, 600, 50)
+    show_code = st.checkbox("Show code")
+
+st.header("Visualization")
+st.text(
+    "A visualization of the famous Cora citation network. Each of its "
+    "seven scientific subjects is represented by a different color."
+)
+
+components.html(
+    render_graph(VG, height=height).data,
+    height=height,
+)
+
+if show_code:
+    st.header("Code")
+    st.code(script_path.read_text())

examples/datasets/cora/cora_rels.parquet.gzip

@@ -47,6 +47,7 @@ dev = [
     "palettable==3.3.3",
     "pytest-mock==3.14.0",
     "nbconvert==7.16.5",
+    "streamlit==1.41.1",
 ]
 pandas = ["pandas>=2, <3", "pandas-stubs>=2, <3"]
 gds = ["graphdatascience>=1, <2"] # not compatible yet with Python 3.13

examples/streamlit-example.py

@@ -2,7 +2,7 @@
 
 from typing import Any, Optional, Union
 
-from pydantic import BaseModel, Field, field_serializer, field_validator
+from pydantic import AliasChoices, BaseModel, Field, field_serializer, field_validator
 from pydantic_extra_types.color import Color, ColorType
 
 from .options import CaptionAlignment
@@ -14,7 +14,9 @@ class Node(BaseModel, extra="allow"):
     All options available in the NVL library (see https://neo4j.com/docs/nvl/current/base-library/#_nodes)
     """
 
-    id: Union[str, int] = Field(description="Unique identifier for the node")
+    id: Union[str, int] = Field(
+        validation_alias=AliasChoices("id", "nodeId", "node_id"), description="Unique identifier for the node"
+    )
     caption: Optional[str] = Field(None, description="The caption of the node")
     caption_align: Optional[CaptionAlignment] = Field(
         None, serialization_alias="captionAlign", description="The alignment of the caption text"

python-wrapper/pyproject.toml

@@ -23,6 +23,11 @@ def __init__(self) -> None:
         with js_path.open("r", encoding="utf-8") as file:
             self.library_code = file.read()
 
+    def unsupported_field_type_error(self, e: TypeError, entity: str) -> Exception:
+        if "not JSON serializable" in str(e):
+            return ValueError(f"A field of a {entity} object is not supported: {str(e)}")
+        return e
+
     def render(
         self,
         nodes: list[Node],
@@ -31,8 +36,14 @@ def render(
         width: str,
         height: str,
     ) -> HTML:
-        nodes_json = json.dumps([node.to_dict() for node in nodes])
-        rels_json = json.dumps([rel.to_dict() for rel in relationships])
+        try:
+            nodes_json = json.dumps([node.to_dict() for node in nodes])
+        except TypeError as e:
+            raise self.unsupported_field_type_error(e, "node")
+        try:
+            rels_json = json.dumps([rel.to_dict() for rel in relationships])
+        except TypeError as e:
+            raise self.unsupported_field_type_error(e, "relationship")
 
         render_options_json = json.dumps(render_options.to_dict())
 

python-wrapper/src/neo4j_viz/node.py

@@ -3,7 +3,7 @@
 from typing import Any, Optional, Union
 from uuid import uuid4
 
-from pydantic import BaseModel, Field, field_serializer, field_validator
+from pydantic import AliasChoices, BaseModel, Field, field_serializer, field_validator
 from pydantic_extra_types.color import Color, ColorType
 
 from .options import CaptionAlignment
@@ -19,9 +19,15 @@ class Relationship(BaseModel, extra="allow"):
         default_factory=lambda: uuid4().hex, description="Unique identifier for the relationship"
     )
     source: Union[str, int] = Field(
-        serialization_alias="from", description="Node ID where the relationship points from"
+        serialization_alias="from",
+        validation_alias=AliasChoices("source", "sourceNodeId", "source_node_id", "from"),
+        description="Node ID where the relationship points from",
+    )
+    target: Union[str, int] = Field(
+        serialization_alias="to",
+        validation_alias=AliasChoices("target", "targetNodeId", "target_node_id", "to"),
+        description="Node ID where the relationship points to",
     )
-    target: Union[str, int] = Field(serialization_alias="to", description="Node ID where the relationship points to")
     caption: Optional[str] = Field(None, description="The caption of the relationship")
     caption_align: Optional[CaptionAlignment] = Field(
         None, serialization_alias="captionAlign", description="The alignment of the caption text"

python-wrapper/src/neo4j_viz/nvl.py

@@ -1,3 +1,5 @@
+import pytest
+
 from neo4j_viz import CaptionAlignment, Node
 
 
@@ -53,3 +55,12 @@ def test_node_with_additional_fields() -> None:
         "id": "1",
         "componentId": 2,
     }
+
+
+@pytest.mark.parametrize("alias", ["id", "nodeId", "node_id"])
+def test_id_aliases(alias: str) -> None:
+    node = Node(**{alias: 1})
+
+    assert node.to_dict() == {
+        "id": "1",
+    }

python-wrapper/src/neo4j_viz/relationship.py

@@ -1,3 +1,5 @@
+import pytest
+
 from neo4j_viz import CaptionAlignment
 from neo4j_viz.relationship import Relationship
 
@@ -47,3 +49,20 @@ def test_rels_additional_fields() -> None:
     rel_dict = rel.to_dict()
     assert {"id", "from", "to", "componentId"} == set(rel_dict.keys())
     assert rel.componentId == 2  # type: ignore[attr-defined]
+
+
+@pytest.mark.parametrize("src_alias", ["source", "sourceNodeId", "source_node_id", "from"])
+@pytest.mark.parametrize("trg_alias", ["target", "targetNodeId", "target_node_id", "to"])
+def test_aliases(src_alias: str, trg_alias: str) -> None:
+    rel = Relationship(
+        **{
+            src_alias: "1",
+            trg_alias: "2",
+        }
+    )
+
+    rel_dict = rel.to_dict()
+
+    assert {"id", "from", "to"} == set(rel_dict.keys())
+    assert rel_dict["from"] == "1"
+    assert rel_dict["to"] == "2"

python-wrapper/tests/test_node.py

@@ -60,3 +60,29 @@ def test_basic_render(render_option: dict[str, Any], tmp_path: Path) -> None:
     severe_logs = [log for log in logs if log["level"] == "SEVERE"]
 
     assert not severe_logs, f"Severe logs found: {severe_logs}, all logs: {logs}"
+
+
+def test_unsupported_field_type() -> None:
+    with pytest.raises(
+        ValueError, match="A field of a node object is not supported: Object of type set is not JSON serializable"
+    ):
+        nodes = [
+            Node(id="4:d09f48a4-5fca-421d-921d-a30a896c604d:0", caption="Person", unsupported=set([1, 2, 3])),
+        ]
+        VG = VisualizationGraph(nodes=nodes, relationships=[])
+        VG.render()
+
+    with pytest.raises(
+        ValueError,
+        match="A field of a relationship object is not supported: Object of type set is not JSON serializable",
+    ):
+        relationships = [
+            Relationship(
+                source="4:d09f48a4-5fca-421d-921d-a30a896c604d:0",
+                target="4:d09f48a4-5fca-421d-921d-a30a896c604d:6",
+                caption="BUYS",
+                unsupported=set([1, 2, 3]),
+            ),
+        ]
+        VG = VisualizationGraph(nodes=[], relationships=relationships)
+        VG.render()