diff --git a/examples/snowpark-example.ipynb b/examples/snowpark-example.ipynb index 51d93348..a1ac48b4 100644 --- a/examples/snowpark-example.ipynb +++ b/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\"})" ] }, { diff --git a/python-wrapper/src/neo4j_viz/node.py b/python-wrapper/src/neo4j_viz/node.py index f8e15c98..93c6036e 100644 --- a/python-wrapper/src/neo4j_viz/node.py +++ b/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 diff --git a/python-wrapper/src/neo4j_viz/nvl.py b/python-wrapper/src/neo4j_viz/nvl.py index f33b5533..e7f9b78c 100644 --- a/python-wrapper/src/neo4j_viz/nvl.py +++ b/python-wrapper/src/neo4j_viz/nvl.py @@ -112,7 +112,7 @@ def render(