Introduce shared GdsBaseModel

Co-authored-by: Florentin Dörre <[email protected]>

Author: DarthMax

graphdatascience/arrow_client/arrow_base_model.py

@@ -1,12 +1,10 @@
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 from pydantic.alias_generators import to_camel
 
 
-class ArrowBaseModel(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class ArrowBaseModel(BaseModel, alias_generator=to_camel):
     def dump_camel(self) -> dict[str, Any]:
         return self.model_dump(by_alias=True)
 

graphdatascience/procedure_surface/api/catalog_endpoints.py

@@ -5,10 +5,10 @@
 from datetime import datetime
 from typing import Any, List, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator
-from pydantic.alias_generators import to_camel
+from pydantic import Field, field_validator
 
 from graphdatascience import Graph
+from graphdatascience.procedure_surface.utils.GdsBaseModel import GdsBaseModel
 
 
 class CatalogEndpoints(ABC):
@@ -66,9 +66,7 @@ def filter(
         pass
 
 
-class GraphListResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class GraphListResult(GdsBaseModel):
     graph_name: str
     database: str
     database_location: str
@@ -91,9 +89,7 @@ def strip_timezone(cls, value: Any) -> Any:
         return value
 
 
-class GraphFilterResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class GraphFilterResult(GdsBaseModel):
     graph_name: str
     from_graph_name: str
     node_filter: str

graphdatascience/procedure_surface/api/estimation_result.py

@@ -2,13 +2,10 @@
 
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+from graphdatascience.procedure_surface.utils.GdsBaseModel import GdsBaseModel
 
 
-class EstimationResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class EstimationResult(GdsBaseModel):
     node_count: int
     relationship_count: int
     required_memory: str
@@ -19,9 +16,6 @@ class EstimationResult(BaseModel):
     heap_percentage_min: float
     heap_percentage_max: float
 
-    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/k1coloring_endpoints.py

@@ -4,10 +4,9 @@
 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 ..utils.GdsBaseModel import GdsBaseModel
 from .estimation_result import EstimationResult
 
 
@@ -245,9 +244,7 @@ def estimate(
         pass
 
 
-class K1ColoringMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class K1ColoringMutateResult(GdsBaseModel):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -257,13 +254,8 @@ class K1ColoringMutateResult(BaseModel):
     mutate_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class K1ColoringStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class K1ColoringStatsResult(GdsBaseModel):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -272,13 +264,8 @@ class K1ColoringStatsResult(BaseModel):
     compute_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class K1ColoringWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class K1ColoringWriteResult(GdsBaseModel):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -287,6 +274,3 @@ class K1ColoringWriteResult(BaseModel):
     compute_millis: int
     write_millis: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)

graphdatascience/procedure_surface/api/kcore_endpoints.py

@@ -4,10 +4,9 @@
 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 ..utils.GdsBaseModel import GdsBaseModel
 from .estimation_result import EstimationResult
 
 
@@ -219,9 +218,7 @@ def estimate(
         pass
 
 
-class KCoreMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class KCoreMutateResult(GdsBaseModel):
     node_properties_written: int
     degeneracy: int
     pre_processing_millis: int
@@ -230,13 +227,8 @@ class KCoreMutateResult(BaseModel):
     mutate_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class KCoreStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class KCoreStatsResult(GdsBaseModel):
     degeneracy: int
     pre_processing_millis: int
     compute_millis: int
@@ -247,16 +239,11 @@ def __getitem__(self, item: str) -> Any:
         return getattr(self, item)
 
 
-class KCoreWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class KCoreWriteResult(GdsBaseModel):
     node_properties_written: int
     degeneracy: int
     pre_processing_millis: int
     compute_millis: int
     post_processing_millis: int
     write_millis: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)

graphdatascience/procedure_surface/api/louvain_endpoints.py

@@ -4,10 +4,9 @@
 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 ..utils.GdsBaseModel import GdsBaseModel
 from .estimation_result import EstimationResult
 
 
@@ -305,9 +304,7 @@ def estimate(
         pass
 
 
-class LouvainMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class LouvainMutateResult(GdsBaseModel):
     modularity: float
     modularities: List[Any]
     ran_levels: int
@@ -320,13 +317,8 @@ class LouvainMutateResult(BaseModel):
     node_properties_written: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class LouvainStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class LouvainStatsResult(GdsBaseModel):
     modularity: float
     modularities: List[Any]
     ran_levels: int
@@ -337,13 +329,8 @@ class LouvainStatsResult(BaseModel):
     post_processing_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class LouvainWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class LouvainWriteResult(GdsBaseModel):
     modularity: float
     modularities: List[Any]
     ran_levels: int
@@ -355,6 +342,3 @@ class LouvainWriteResult(BaseModel):
     write_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)

graphdatascience/procedure_surface/api/scc_endpoints.py

@@ -4,10 +4,9 @@
 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 ..utils.GdsBaseModel import GdsBaseModel
 from .estimation_result import EstimationResult
 
 
@@ -227,9 +226,7 @@ def estimate(
         pass
 
 
-class SccMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class SccMutateResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -239,27 +236,17 @@ class SccMutateResult(BaseModel):
     node_properties_written: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class SccStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class SccStatsResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
     compute_millis: int
     post_processing_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class SccWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class SccWriteResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -268,6 +255,3 @@ class SccWriteResult(BaseModel):
     post_processing_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)

graphdatascience/procedure_surface/api/wcc_endpoints.py

@@ -4,10 +4,9 @@
 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 ..utils.GdsBaseModel import GdsBaseModel
 from .estimation_result import EstimationResult
 
 
@@ -267,9 +266,7 @@ def estimate(
         pass
 
 
-class WccMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class WccMutateResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -279,27 +276,17 @@ class WccMutateResult(BaseModel):
     node_properties_written: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class WccStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class WccStatsResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
     compute_millis: int
     post_processing_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)
-
-
-class WccWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class WccWriteResult(GdsBaseModel):
     component_count: int
     component_distribution: dict[str, Any]
     pre_processing_millis: int
@@ -308,6 +295,3 @@ class WccWriteResult(BaseModel):
     post_processing_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return getattr(self, item)