Merge pull request #165 from neo4j/case-insensitive-continued

Consider all field aliases in convenience constructors

Author: adamnsch

examples/snowpark-example.ipynb

@@ -249,13 +249,8 @@ def from_gql_create(
     node_pattern = re.compile(r"^\(([^)]*)\)$")
     rel_pattern = re.compile(r"^\(([^)]*)\)-\s*\[\s*:(\w+)\s*(\{[^}]*\})?\s*\]->\(([^)]*)\)$")
 
-    node_top_level_keys = set(Node.model_fields.keys())
-    node_top_level_keys.remove("id")
-
-    rel_top_level_keys = set(Relationship.model_fields.keys())
-    rel_top_level_keys.remove("id")
-    rel_top_level_keys.remove("source")
-    rel_top_level_keys.remove("target")
+    node_top_level_keys = Node.all_validation_aliases(exempted_fields=["id"])
+    rel_top_level_keys = Relationship.all_validation_aliases(exempted_fields=["id", "source", "target"])
 
     nodes = []
     relationships = []

python-wrapper/src/neo4j_viz/gql_create.py

@@ -47,10 +47,15 @@ def from_neo4j(
     else:
         raise ValueError(f"Invalid input type `{type(result)}`. Expected `neo4j.Graph` or `neo4j.Result`")
 
-    nodes = [_map_node(node, size_property, caption_property=node_caption) for node in graph.nodes]
+    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
+    ]
     relationships = []
     for rel in graph.relationships:
-        mapped_rel = _map_relationship(rel, caption_property=relationship_caption)
+        mapped_rel = _map_relationship(rel, all_rel_field_aliases, caption_property=relationship_caption)
         if mapped_rel:
             relationships.append(mapped_rel)
 
@@ -62,7 +67,12 @@ def from_neo4j(
     return VG
 
 
-def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_property: Optional[str]) -> Node:
+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:
@@ -78,7 +88,7 @@ def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_prop
 
     properties = {}
     for prop, value in node.items():
-        if prop not in Node.model_fields.keys():
+        if prop not in all_node_field_aliases:
             properties[prop] = value
             continue
 
@@ -95,7 +105,9 @@ def _map_node(node: neo4j.graph.Node, size_property: Optional[str], caption_prop
     return Node(**top_level_fields, properties=properties)
 
 
-def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[str]) -> Optional[Relationship]:
+def _map_relationship(
+    rel: neo4j.graph.Relationship, all_rel_field_aliases: set[str], caption_property: Optional[str]
+) -> Optional[Relationship]:
     if rel.start_node is None or rel.end_node is None:
         return None
 
@@ -109,7 +121,7 @@ def _map_relationship(rel: neo4j.graph.Relationship, caption_property: Optional[
 
     properties = {}
     for prop, value in rel.items():
-        if prop not in Relationship.model_fields.keys():
+        if prop not in all_rel_field_aliases:
             properties[prop] = value
             continue
 

python-wrapper/src/neo4j_viz/neo4j.py

@@ -88,3 +88,12 @@ def cast_color(cls, color: ColorType) -> Color:
 
     def to_dict(self) -> dict[str, Any]:
         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
+
+        return {str(alias) for aliases in by_field for alias in aliases}

python-wrapper/src/neo4j_viz/node.py

@@ -23,6 +23,9 @@ def _from_dfs(
     else:
         node_dfs_iter = node_dfs
 
+    all_node_field_aliases = Node.all_validation_aliases()
+    all_rel_field_aliases = Relationship.all_validation_aliases()
+
     has_size = True
     nodes = []
     for node_df in node_dfs_iter:
@@ -31,7 +34,7 @@ def _from_dfs(
             top_level = {}
             properties = {}
             for key, value in row.to_dict().items():
-                if key in Node.model_fields.keys():
+                if key in all_node_field_aliases:
                     top_level[key] = value
                 else:
                     if rename_properties and key in rename_properties:
@@ -51,7 +54,7 @@ def _from_dfs(
             top_level = {}
             properties = {}
             for key, value in row.to_dict().items():
-                if key in Relationship.model_fields.keys():
+                if key in all_rel_field_aliases:
                     top_level[key] = value
                 else:
                     if rename_properties and key in rename_properties:

python-wrapper/src/neo4j_viz/pandas.py

@@ -95,3 +95,16 @@ def cast_color(cls, color: ColorType) -> Color:
 
     def to_dict(self) -> dict[str, Any]:
         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  # type: ignore
+            for k, v in Relationship.model_fields.items()
+            if k not in exempted_fields
+        ]
+
+        return {str(alias) for aliases in by_field for alias in aliases}

python-wrapper/src/neo4j_viz/relationship.py

@@ -93,3 +93,15 @@ def test_node_casing() -> None:
     assert node.caption == "Person"
     assert node.caption_align == CaptionAlignment.TOP
     assert node.caption_size == 1
+
+
+def test_all_validation_aliases() -> None:
+    all_aliases = Node.all_validation_aliases()
+    assert "CAPTION_ALIGN" in all_aliases
+    assert "captionAlign" in all_aliases
+    assert "caption_align" in all_aliases
+
+    all_aliases = Node.all_validation_aliases(exempted_fields=["caption_align"])
+    assert "CAPTION_ALIGN" not in all_aliases
+    assert "captionAlign" not in all_aliases
+    assert "caption_align" not in all_aliases

python-wrapper/tests/test_node.py

@@ -101,3 +101,15 @@ def test_rel_casing() -> None:
     assert rel.id == "1"
     assert rel.caption_align == CaptionAlignment.TOP
     assert rel.caption_size == 12
+
+
+def test_all_validation_aliases() -> None:
+    all_aliases = Relationship.all_validation_aliases()
+    assert "CAPTION_ALIGN" in all_aliases
+    assert "captionAlign" in all_aliases
+    assert "caption_align" in all_aliases
+
+    all_aliases = Relationship.all_validation_aliases(exempted_fields=["caption_align"])
+    assert "CAPTION_ALIGN" not in all_aliases
+    assert "captionAlign" not in all_aliases
+    assert "caption_align" not in all_aliases