Use BaseResult consistently

Author: FlorentinD

graphdatascience/procedure_surface/api/articlerank_endpoints.py

@@ -4,8 +4,8 @@
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -304,9 +304,7 @@ def estimate(
         """
 
 
-class ArticleRankMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class ArticleRankMutateResult(BaseResult):
     ran_iterations: int
     did_converge: bool
     centrality_distribution: dict[str, Any]
@@ -317,13 +315,8 @@ class ArticleRankMutateResult(BaseModel):
     node_properties_written: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
-
-
-class ArticleRankStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class ArticleRankStatsResult(BaseResult):
     ran_iterations: int
     did_converge: bool
     centrality_distribution: dict[str, Any]
@@ -332,13 +325,8 @@ class ArticleRankStatsResult(BaseModel):
     post_processing_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
-
-
-class ArticleRankWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class ArticleRankWriteResult(BaseResult):
     ran_iterations: int
     did_converge: bool
     centrality_distribution: dict[str, Any]
@@ -348,6 +336,3 @@ class ArticleRankWriteResult(BaseModel):
     write_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]

graphdatascience/procedure_surface/api/articulationpoints_endpoints.py

@@ -2,8 +2,8 @@
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -215,44 +215,29 @@ def estimate(
         pass
 
 
-class ArticulationPointsMutateResult(BaseModel):
+class ArticulationPointsMutateResult(BaseResult):
     """Result of running ArticulationPoints algorithm with mutate mode."""
 
-    model_config = ConfigDict(alias_generator=to_camel)
-
     articulation_point_count: int
     node_properties_written: int
     mutate_millis: int
     compute_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
 
-
-class ArticulationPointsStatsResult(BaseModel):
+class ArticulationPointsStatsResult(BaseResult):
     """Result of running ArticulationPoints algorithm with stats mode."""
 
-    model_config = ConfigDict(alias_generator=to_camel)
-
     articulation_point_count: int
     compute_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
 
-
-class ArticulationPointsWriteResult(BaseModel):
+class ArticulationPointsWriteResult(BaseResult):
     """Result of running ArticulationPoints algorithm with write mode."""
 
-    model_config = ConfigDict(alias_generator=to_camel)
-
     articulation_point_count: int
     node_properties_written: int
     write_millis: int
     compute_millis: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]

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.api.base_result import BaseResult
 
 
-class EstimationResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class EstimationResult(BaseResult):
     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/fastrp_endpoints.py

@@ -4,8 +4,8 @@
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -337,38 +337,23 @@ def estimate(
         """
 
 
-class FastRPMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class FastRPMutateResult(BaseResult):
     pre_processing_millis: int
     compute_millis: int
     mutate_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
-
-
-class FastRPStatsResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class FastRPStatsResult(BaseResult):
     pre_processing_millis: int
     compute_millis: int
     configuration: dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
-
-
-class FastRPWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class FastRPWriteResult(BaseResult):
     pre_processing_millis: int
     compute_millis: int
     write_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]

graphdatascience/procedure_surface/api/hashgnn_endpoints.py

@@ -4,8 +4,8 @@
 from typing import Any, Dict, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -245,37 +245,27 @@ def estimate(
         """
 
 
-class HashGNNMutateResult(BaseModel):
+class HashGNNMutateResult(BaseResult):
     """
     Result object representing the results of running a HashGNN algorithm in mutate mode.
     """
 
-    model_config = ConfigDict(alias_generator=to_camel)
-
     node_count: int
     node_properties_written: int
     pre_processing_millis: int
     compute_millis: int
     mutate_millis: int
     configuration: Dict[str, Any]
 
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]
-
 
-class HashGNNWriteResult(BaseModel):
+class HashGNNWriteResult(BaseResult):
     """
     Result object representing the results of running a HashGNN algorithm in write mode.
     """
 
-    model_config = ConfigDict(alias_generator=to_camel)
-
     node_count: int
     node_properties_written: int
     pre_processing_millis: int
     compute_millis: int
     write_millis: int
     configuration: Dict[str, Any]
-
-    def __getitem__(self, item: str) -> Any:
-        return self.__dict__[item]

graphdatascience/procedure_surface/api/k1coloring_endpoints.py

@@ -4,8 +4,8 @@
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -255,9 +255,7 @@ def estimate(
         pass
 
 
-class K1ColoringMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class K1ColoringMutateResult(BaseResult):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -267,13 +265,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(BaseResult):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -282,13 +275,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(BaseResult):
     node_count: int
     color_count: int
     ran_iterations: int
@@ -297,6 +285,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,8 +4,8 @@
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict
-from pydantic.alias_generators import to_camel
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
 
 from ...graph.graph_object import Graph
 from .estimation_result import EstimationResult
@@ -225,9 +225,7 @@ def estimate(
         pass
 
 
-class KCoreMutateResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
-
+class KCoreMutateResult(BaseResult):
     node_properties_written: int
     degeneracy: int
     pre_processing_millis: int
@@ -236,33 +234,20 @@ 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(BaseResult):
     degeneracy: int
     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 KCoreWriteResult(BaseModel):
-    model_config = ConfigDict(alias_generator=to_camel)
 
+class KCoreWriteResult(BaseResult):
     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)