Return custom type from wcc.mutate

Author: DarthMax

graphdatascience/procedure_surface/api/wcc_endpoints.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Any, List, Optional
 
 from pandas import DataFrame, Series
@@ -27,7 +30,7 @@ def mutate(
         seed_property: Optional[str] = None,
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
-    ) -> Series[Any]:
+    ) -> WccMutateResult:
         """
         Executes the WCC algorithm and writes the results to the in-memory graph as node properties.
 
@@ -82,7 +85,7 @@ def stats(
         seed_property: Optional[str] = None,
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
-    ) -> Series[Any]:
+    ) -> Series:
         """
         Executes the WCC algorithm and returns statistics.
 
@@ -195,7 +198,7 @@ def write(
         relationship_weight_property: Optional[str] = None,
         write_concurrency: Optional[Any] = None,
         write_to_result_store: Optional[bool] = None,
-    ) -> Series[Any]:
+    ) -> Series:
         """
         Executes the WCC algorithm and writes the results to the Neo4j database.
 
@@ -240,3 +243,15 @@ def write(
             Algorithm metrics and statistics
         """
         pass
+
+
+@dataclass(frozen=True, repr=True)
+class WccMutateResult:
+    component_count: int
+    component_distribution: dict[str, Any]
+    pre_processing_millis: int
+    compute_millis: int
+    post_processing_millis: int
+    mutate_millis: int
+    node_properties_written: int
+    configuration: dict[str, Any]

graphdatascience/procedure_surface/cypher/wcc_proc_runner.py

@@ -5,7 +5,7 @@
 from ...call_parameters import CallParameters
 from ...graph.graph_object import Graph
 from ...query_runner.query_runner import QueryRunner
-from ..api.wcc_endpoints import WccEndpoints
+from ..api.wcc_endpoints import WccEndpoints, WccMutateResult
 
 
 class WccCypherEndpoints(WccEndpoints):
@@ -32,7 +32,7 @@ def mutate(
         seed_property: Optional[str] = None,
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
-    ) -> Series[Any]:
+    ) -> WccMutateResult:
         # Build configuration dictionary from parameters
         config: dict[str, Any] = {
             "mutateProperty": mutate_property,
@@ -66,7 +66,17 @@ def mutate(
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
-        return self._query_runner.call_procedure(endpoint="gds.wcc.mutate", params=params).squeeze()  # type: ignore
+        cypher_result = self._query_runner.call_procedure(endpoint="gds.wcc.mutate", params=params).squeeze()
+
+        return WccMutateResult(
+            cypher_result["componentCount"],
+            cypher_result["componentDistribution"],
+            cypher_result["preProcessingMillis"],
+            cypher_result["computeMillis"],
+            cypher_result["postProcessingMillis"],
+            cypher_result["mutateMillis"],
+            cypher_result["nodePropertiesWritten"],
+        )
 
     def stats(
         self,
@@ -82,7 +92,7 @@ def stats(
         seed_property: Optional[str] = None,
         consecutive_ids: Optional[bool] = None,
         relationship_weight_property: Optional[str] = None,
-    ) -> Series[Any]:
+    ) -> Series:
         # Build configuration dictionary from parameters
         config: dict[str, Any] = {}
 
@@ -185,7 +195,7 @@ def write(
         relationship_weight_property: Optional[str] = None,
         write_concurrency: Optional[int] = None,
         write_to_result_store: Optional[bool] = None,
-    ) -> Series[Any]:
+    ) -> Series:
         # Build configuration dictionary from parameters
         config: dict[str, Any] = {
             "writeProperty": write_property,