Use pydantic to parse cypher results into result classes

Author: DarthMax

graphdatascience/procedure_surface/api/estimation_result.py

@@ -1,9 +1,14 @@
-from dataclasses import dataclass
+from __future__ import annotations
+
 from typing import Any
 
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class EstimationResult(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
 
-@dataclass(frozen=True, repr=True)
-class EstimationResult:
     node_count: int
     relationship_count: int
     required_memory: str
@@ -16,3 +21,7 @@ class EstimationResult:
 
     def __getitem__(self, item: str) -> Any:
         return getattr(self, item)
+
+    @staticmethod
+    def from_cypher(cypher_result: dict[str, Any]) -> EstimationResult:
+        return EstimationResult(**cypher_result)

graphdatascience/procedure_surface/api/wcc_endpoints.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from typing import Any, List, Optional
 
 from pandas import DataFrame
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -272,8 +273,9 @@ def estimate(
         pass
 
 
-@dataclass(frozen=True, repr=True)
-class WccMutateResult:
+class WccMutateResult(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
+
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -287,8 +289,9 @@ def __getitem__(self, item: str) -> Any:
         return getattr(self, item)
 
 
-@dataclass(frozen=True, repr=True)
-class WccStatsResult:
+class WccStatsResult(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
+
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -300,8 +303,9 @@ def __getitem__(self, item: str) -> Any:
         return getattr(self, item)
 
 
-@dataclass(frozen=True, repr=True)
-class WccWriteResult:
+class WccWriteResult(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
+
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int

requirements/base/base.txt

@@ -8,3 +8,4 @@ tqdm >= 4.0, < 5.0
 typing-extensions >= 4.0, < 5.0
 requests
 tenacity >= 9.0
+pydantic >= 2.11