Make Node and Relationship top-level fields case insensitive

examples/snowpark-example.ipynb

@@ -171,7 +171,7 @@
     "## Fetching the data\n",
     "\n",
     "Next we fetch our tables from Snowflake and convert them to pandas DataFrames.\n",
-    "Additionally, we rename the most of the table columns so that they are named according to the `neo4j-viz` API."
+    "Additionally, we rename some of the table columns so that they are named according to the `neo4j-viz` API."
    ]
   },
   {
@@ -181,16 +181,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "products_df = (\n",
-    "    session.table(\"products\")\n",
-    "    .to_pandas()\n",
-    "    .rename(columns={\"ID\": \"id\", \"NAME\": \"caption\"})\n",
-    ")\n",
-    "parents_df = (\n",
-    "    session.table(\"parents\")\n",
-    "    .to_pandas()\n",
-    "    .rename(columns={\"SOURCE\": \"source\", \"TARGET\": \"target\", \"TYPE\": \"caption\"})\n",
-    ")"
+    "products_df = session.table(\"products\").to_pandas().rename(columns={\"NAME\": \"caption\"})\n",
+    "parents_df = session.table(\"parents\").to_pandas().rename(columns={\"TYPE\": \"caption\"})"
    ]
   },
   {

python-wrapper/src/neo4j_viz/node.py

@@ -2,7 +2,8 @@
 
 from typing import Any, Optional, Union
 
-from pydantic import AliasChoices, BaseModel, Field, field_serializer, field_validator
+from pydantic import AliasChoices, AliasGenerator, BaseModel, Field, field_serializer, field_validator
+from pydantic.alias_generators import to_camel
 from pydantic_extra_types.color import Color, ColorType
 
 from .node_size import RealNumber
@@ -11,29 +12,46 @@
 NodeIdType = Union[str, int]
 
 
-class Node(BaseModel, extra="allow"):
+def create_aliases(field_name: str) -> AliasChoices:
+    valid_names = [field_name]
+
+    if field_name == "id":
+        valid_names.extend(["nodeid", "node_id"])
+
+    choices = [[choice, choice.upper(), to_camel(choice)] for choice in valid_names]
+
+    return AliasChoices(*[alias for aliases in choices for alias in aliases])
+
+
+class Node(
+    BaseModel,
+    extra="forbid",
+    alias_generator=AliasGenerator(
+        validation_alias=create_aliases,
+        serialization_alias=lambda field_name: to_camel(field_name),
+    ),
+):
     """
     A node in a graph to visualize.
 
+    Each field is case-insensitive for input, and camelCase is also accepted.
+    For example, "CAPTION_ALIGN", "captionAlign" are also valid inputs keys for the `caption_align` field.
+    Upon construction however, the field names are converted to snake_case.
+
     For more info on each field, see the NVL library docs: https://neo4j.com/docs/nvl/current/base-library/#_nodes
     """
 
     #: Unique identifier for the node
-    id: NodeIdType = Field(
-        validation_alias=AliasChoices("id", "nodeId", "node_id"), description="Unique identifier for the node"
-    )
+    id: NodeIdType = Field(description="Unique identifier for the node")
     #: The caption of the node
     caption: Optional[str] = Field(None, description="The caption of the node")
     #: The alignment of the caption text
-    caption_align: Optional[CaptionAlignment] = Field(
-        None, serialization_alias="captionAlign", description="The alignment of the caption text"
-    )
+    caption_align: Optional[CaptionAlignment] = Field(None, description="The alignment of the caption text")
     #: The size of the caption text. The font size to node radius ratio
     caption_size: Optional[int] = Field(
         None,
         ge=1,
         le=3,
-        serialization_alias="captionSize",
         description="The size of the caption text. The font size to node radius ratio",
     )
     #: The size of the node as radius in pixel

python-wrapper/src/neo4j_viz/nvl.py

@@ -112,7 +112,7 @@ def render(
         <script>
             getTheme = () => {{
                 const backgroundColorString = window.getComputedStyle(document.body, null).getPropertyValue('background-color')
-                const colorsArray = backgroundColorString.match(/\d+/g);
+                const colorsArray = backgroundColorString.match(/\\d+/g);
                 const brightness = Number(colorsArray[0]) * 0.2126 + Number(colorsArray[1]) * 0.7152 + Number(colorsArray[2]) * 0.0722
                 return brightness < 128 ? "dark" : "light"
             }}

python-wrapper/src/neo4j_viz/relationship.py

@@ -3,16 +3,41 @@
 from typing import Any, Optional, Union
 from uuid import uuid4
 
-from pydantic import AliasChoices, BaseModel, Field, field_serializer, field_validator
+from pydantic import AliasChoices, AliasGenerator, BaseModel, Field, field_serializer, field_validator
+from pydantic.alias_generators import to_camel
 from pydantic_extra_types.color import Color, ColorType
 
 from .options import CaptionAlignment
 
 
-class Relationship(BaseModel, extra="allow"):
+def create_aliases(field_name: str) -> AliasChoices:
+    valid_names = [field_name]
+
+    if field_name == "source":
+        valid_names.extend(["sourcenodeid", "source_node_id", "from"])
+    if field_name == "target":
+        valid_names.extend(["targetnodeid", "target_node_id", "to"])
+
+    choices = [[choice, choice.upper(), to_camel(choice)] for choice in valid_names]
+
+    return AliasChoices(*[alias for aliases in choices for alias in aliases])
+
+
+class Relationship(
+    BaseModel,
+    extra="forbid",
+    alias_generator=AliasGenerator(
+        validation_alias=create_aliases,
+        serialization_alias=lambda field_name: to_camel(field_name),
+    ),
+):
     """
     A relationship in a graph to visualize.
 
+    Each field is case-insensitive for input, and camelCase is also accepted.
+    For example, "CAPTION_ALIGN", "captionAlign" are also valid inputs keys for the `caption_align` field.
+    Upon construction however, the field names are converted to snake_case.
+
     For more info on each field, see the NVL library docs: https://neo4j.com/docs/nvl/current/base-library/#_relationships
     """
 
@@ -23,25 +48,19 @@ class Relationship(BaseModel, extra="allow"):
     #: Node ID where the relationship points from
     source: Union[str, int] = Field(
         serialization_alias="from",
-        validation_alias=AliasChoices("source", "sourceNodeId", "source_node_id", "from"),
         description="Node ID where the relationship points from",
     )
     #: Node ID where the relationship points to
     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",
     )
     #: The caption of the relationship
     caption: Optional[str] = Field(None, description="The caption of the relationship")
     #: The alignment of the caption text
-    caption_align: Optional[CaptionAlignment] = Field(
-        None, serialization_alias="captionAlign", description="The alignment of the caption text"
-    )
+    caption_align: Optional[CaptionAlignment] = Field(None, description="The alignment of the caption text")
     #: The size of the caption text
-    caption_size: Optional[Union[int, float]] = Field(
-        None, gt=0.0, serialization_alias="captionSize", description="The size of the caption text"
-    )
+    caption_size: Optional[Union[int, float]] = Field(None, gt=0.0, description="The size of the caption text")
     #: The color of the relationship. Allowed input is for example "#FF0000", "red" or (255, 0, 0)
     color: Optional[ColorType] = Field(None, description="The color of the relationship")
     #: Additional properties of the relationship that do not directly impact the visualization

python-wrapper/tests/test_node.py

@@ -54,20 +54,19 @@ def test_node_with_float_size() -> None:
     }
 
 
-def test_node_with_additional_fields() -> None:
+def test_node_with_properties() -> None:
     node = Node(
         id="1",
-        componentId=2,
+        properties=dict(componentId=2),
     )
 
     assert node.to_dict() == {
         "id": "1",
-        "componentId": 2,
-        "properties": {},
+        "properties": {"componentId": 2},
     }
 
 
-@pytest.mark.parametrize("alias", ["id", "nodeId", "node_id"])
+@pytest.mark.parametrize("alias", ["id", "nodeId", "node_id", "NODEID", "nodeid"])
 def test_id_aliases(alias: str) -> None:
     node = Node(**{alias: 1})
 
@@ -80,3 +79,17 @@ def test_id_aliases(alias: str) -> None:
 def test_node_validation() -> None:
     with pytest.raises(ValueError, match="Input should be a valid integer, unable to parse string as an integer"):
         Node(id="1", x="not a number")
+
+
+def test_node_casing() -> None:
+    node = Node(
+        ID="4:d09f48a4-5fca-421d-921d-a30a896c604d:0",
+        caption="Person",
+        captionAlign=CaptionAlignment.TOP,
+        CAPTION_SIZE=1,
+    )
+
+    assert node.id == "4:d09f48a4-5fca-421d-921d-a30a896c604d:0"
+    assert node.caption == "Person"
+    assert node.caption_align == CaptionAlignment.TOP
+    assert node.caption_size == 1

python-wrapper/tests/test_relationship.py

@@ -52,12 +52,32 @@ def test_rels_additional_fields() -> None:
     assert rel.properties["componentId"] == 2
 
 
-@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:
+@pytest.mark.parametrize(
+    "src_alias",
+    ["source", "sourceNodeId", "source_node_id", "from", "SOURCE", "SOURCE_NODE_ID", "SOURCENODEID", "FROM"],
+)
+def test_src_aliases(src_alias: str) -> None:
     rel = Relationship(
         **{
             src_alias: "1",
+            "to": "2",
+        }
+    )
+
+    rel_dict = rel.to_dict()
+
+    assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
+    assert rel_dict["from"] == "1"
+    assert rel_dict["to"] == "2"
+
+
+@pytest.mark.parametrize(
+    "trg_alias", ["target", "targetNodeId", "target_node_id", "to", "TARGET", "TARGET_NODE_ID", "TARGETNODEID", "TO"]
+)
+def test_trg_aliases(trg_alias: str) -> None:
+    rel = Relationship(
+        **{
+            "from": "1",
             trg_alias: "2",
         }
     )
@@ -67,3 +87,17 @@ def test_aliases(src_alias: str, trg_alias: str) -> None:
     assert {"id", "from", "to", "properties"} == set(rel_dict.keys())
     assert rel_dict["from"] == "1"
     assert rel_dict["to"] == "2"
+
+
+def test_rel_casing() -> None:
+    rel = Relationship(
+        ID="1",
+        source="2",
+        target="3",
+        captionAlign=CaptionAlignment.TOP,
+        CAPTION_SIZE=12,
+    )
+
+    assert rel.id == "1"
+    assert rel.caption_align == CaptionAlignment.TOP
+    assert rel.caption_size == 12