Refactor flatten_resolve_response (#228)

Move the logic from a standalone `flatten_respolve_response` function,
to a method of the Response object. That way it complements `to_dict`,
`to_json` with `to_flat_dict`.

Author: jm-rivera

datacommons_client/endpoints/resolve.py

@@ -1,35 +1,11 @@
-from typing import Any, Optional
+from typing import Optional
 
 from datacommons_client.endpoints.base import API
 from datacommons_client.endpoints.base import Endpoint
 from datacommons_client.endpoints.payloads import ResolveRequestPayload
 from datacommons_client.endpoints.response import ResolveResponse
 
 
-def flatten_resolve_response(
-    data: ResolveResponse) -> dict[str, list[str] | str]:
-  """
-    Flattens resolved candidate data into a dictionary where each node maps to its candidates.
-
-    Args:
-        data (ResolveResponse): The response object containing the resolved data.
-
-    Returns:
-        dict[str, Any]: A dictionary mapping nodes to their candidates.
-        If a node has only one candidate, it maps directly to the candidate instead of a list.
-    """
-  items: dict[str, Any] = {}
-
-  for entity in data.entities:
-    node = entity.node
-    if len(entity.candidates) == 1:
-      items[node] = entity.candidates[0].dcid
-    else:
-      items[node] = [candidate.dcid for candidate in entity.candidates]
-
-  return items
-
-
 def _resolve_correspondence_expression(from_type: str,
                                        to_type: str,
                                        entity_type: str | None = None) -> str:

datacommons_client/endpoints/response.py

@@ -250,3 +250,22 @@ def from_json(cls, json_data: Dict[str, Any]) -> "ResolveResponse":
     return cls(entities=[
         Entity.from_json(entity) for entity in json_data.get("entities", [])
     ])
+
+  def to_flat_dict(self) -> dict[str, list[str] | str]:
+    """
+      Flattens resolved candidate data into a dictionary where each node maps to its candidates.
+
+      Returns:
+          dict[str, Any]: A dictionary mapping nodes to their candidates.
+          If a node has only one candidate, it maps directly to the candidate instead of a list.
+      """
+    items: dict[str, Any] = {}
+
+    for entity in self.entities:
+      node = entity.node
+      if len(entity.candidates) == 1:
+        items[node] = entity.candidates[0].dcid
+      else:
+        items[node] = [candidate.dcid for candidate in entity.candidates]
+
+    return items

datacommons_client/tests/endpoints/test_resolve_endpoint.py

@@ -2,7 +2,6 @@
 
 from datacommons_client.endpoints.base import API
 from datacommons_client.endpoints.resolve import _resolve_correspondence_expression
-from datacommons_client.endpoints.resolve import flatten_resolve_response
 from datacommons_client.endpoints.resolve import ResolveEndpoint
 from datacommons_client.endpoints.response import ResolveResponse
 
@@ -119,7 +118,7 @@ def test_flatten_resolve_response():
   ])
 
   # Call the function
-  result = flatten_resolve_response(mock_data)
+  result = mock_data.to_flat_dict()
 
   # Expected output
   expected = {