Fix DOT sanitized key collisions in importer

Co-authored-by: Dr Matt James <mattkjames7@users.noreply.github.com>

Author: cursoragent

gqlalchemy/transformations/importing/graph_importer.py

@@ -157,10 +157,7 @@ def _walk(current_graph) -> None:
                 if node_id in {"", "node", "edge", "graph"}:
                     continue
 
-                properties = {
-                    self._sanitize_property_key(key): self._normalize_dot_value(value)
-                    for key, value in (node.get_attributes() or {}).items()
-                }
+                properties = self._normalize_dot_attributes(node.get_attributes() or {})
                 if node_id in graph.nodes:
                     graph.nodes[node_id].update(properties)
                 else:
@@ -172,10 +169,7 @@ def _walk(current_graph) -> None:
                 if not source or not dest:
                     continue
 
-                properties = {
-                    self._sanitize_property_key(key): self._normalize_dot_value(value)
-                    for key, value in (edge.get_attributes() or {}).items()
-                }
+                properties = self._normalize_dot_attributes(edge.get_attributes() or {})
                 graph.add_edge(source, dest, **properties)
 
             for subgraph in current_graph.get_subgraphs():
@@ -186,9 +180,7 @@ def _walk(current_graph) -> None:
         return graph
 
     def _normalize_dot_properties(self, properties: Dict[str, Any]) -> Dict[str, Any]:
-        normalized_attributes = {
-            self._sanitize_property_key(key): self._normalize_dot_value(value) for key, value in properties.items()
-        }
+        normalized_attributes = self._normalize_dot_attributes(properties)
         normalized_properties: Dict[str, Any] = dict(normalized_attributes)
         normalized_properties["attributes_json"] = json.dumps(normalized_attributes, sort_keys=True)
 
@@ -200,6 +192,23 @@ def _normalize_dot_properties(self, properties: Dict[str, Any]) -> Dict[str, Any
 
         return normalized_properties
 
+    def _normalize_dot_attributes(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
+        normalized_attributes: Dict[str, Any] = {}
+        for key, value in attributes.items():
+            sanitized_key = self._sanitize_property_key(key)
+            unique_key = self._resolve_key_collision(sanitized_key, normalized_attributes)
+            normalized_attributes[unique_key] = self._normalize_dot_value(value)
+        return normalized_attributes
+
+    @staticmethod
+    def _resolve_key_collision(base_key: str, properties: Dict[str, Any]) -> str:
+        unique_key = base_key
+        suffix = 1
+        while unique_key in properties:
+            unique_key = f"{base_key}_{suffix}"
+            suffix += 1
+        return unique_key
+
     @staticmethod
     def _normalize_dot_value(value: Any) -> Any:
         if isinstance(value, str):

tests/transformations/importing/test_import.py

@@ -57,6 +57,76 @@ def test_import_nx_normalize_dot_properties_json_key_collision():
     assert normalized_properties["attributes_label"] == "A"
 
 
+def test_import_nx_normalize_dot_properties_sanitized_key_collision_preserves_all_values():
+    importer = GraphImporter(graph_type="Nx")
+
+    normalized_properties = importer._normalize_dot_properties({"my.attr": "dot-value", "my-attr": "dash-value"})
+
+    assert json.loads(normalized_properties["attributes_json"]) == {"my_attr": "dot-value", "my_attr_1": "dash-value"}
+    assert normalized_properties["attributes_my_attr"] == "dot-value"
+    assert normalized_properties["attributes_my_attr_1"] == "dash-value"
+
+
+def test_import_nx_graph_from_pydot_sanitized_key_collision_preserves_all_values():
+    importer = GraphImporter(graph_type="Nx")
+
+    class _DummyNode:
+        def __init__(self, name, attributes):
+            self._name = name
+            self._attributes = attributes
+
+        def get_name(self):
+            return self._name
+
+        def get_attributes(self):
+            return self._attributes
+
+    class _DummyEdge:
+        def __init__(self, source, destination, attributes):
+            self._source = source
+            self._destination = destination
+            self._attributes = attributes
+
+        def get_source(self):
+            return self._source
+
+        def get_destination(self):
+            return self._destination
+
+        def get_attributes(self):
+            return self._attributes
+
+    class _DummyDotGraph:
+        def __init__(self, nodes, edges):
+            self._nodes = nodes
+            self._edges = edges
+
+        def get_nodes(self):
+            return self._nodes
+
+        def get_edges(self):
+            return self._edges
+
+        def get_subgraphs(self):
+            return []
+
+    dot_graph = _DummyDotGraph(
+        nodes=[
+            _DummyNode("A", {"my.attr": "dot-value", "my-attr": "dash-value"}),
+            _DummyNode("B", {}),
+        ],
+        edges=[_DummyEdge("A", "B", {"edge.attr": "left", "edge-attr": "right"})],
+    )
+
+    graph = importer._graph_from_pydot(dot_graph)
+    edge_data = next(iter(graph.edges(data=True)))[2]
+
+    assert graph.nodes["A"]["my_attr"] == "dot-value"
+    assert graph.nodes["A"]["my_attr_1"] == "dash-value"
+    assert edge_data["edge_attr"] == "left"
+    assert edge_data["edge_attr_1"] == "right"
+
+
 def test_import_nx_normalize_dot_graph_for_digraph():
     importer = GraphImporter(graph_type="Nx")
     dot_graph = nx.DiGraph()