Improve response output (#220)

This PR introduces a way to remove `None` values from serialised
responses so they are more compact (and so that they mirror the REST API
response). Additionally, it changes how the data is returned to improve
clarity for users (by using clearer/more appropriate names):
- Remove the `.json` property. The rationale is that we weren't actually
returning a json string but a dictionary. Also, given that it was a
property, we couldn't take additional parameters to specify whether a
'compact' version should be returned.
- Add a `.to_dict()` method, which returns a dictionary. It takes a
parameter `exclude_none` (which defaults to `True`) which excludes any
keys which are None or empty lists*. I decided to call it `exclude_none`
rather than `compact` to mirror how Pydantic deals with the same issue
(model_dump optionally takes `exclude_none=True` in order to remove all
`None` values). This prepares us for a future move to Pydantic (instead
of Python dataclasses), as we've discussed.
- Add a `.to_json()` method, which returns a json string (via
`json.dumps`). It takes the same `exclude_none` parameter described
above.

Those two methods are implemented via a mixin class.

Author: jm-rivera

datacommons_client/endpoints/response.py

@@ -1,6 +1,7 @@
 from dataclasses import asdict
 from dataclasses import dataclass
 from dataclasses import field
+import json
 from typing import Any, Dict, List
 
 from datacommons_client.models.node import Arcs
@@ -14,11 +15,46 @@
 from datacommons_client.models.resolve import Entity
 from datacommons_client.utils.data_processing import flatten_properties
 from datacommons_client.utils.data_processing import observations_as_records
-from datacommons_client.utils.data_processing import unpack_arcs
+
+
+class SerializableMixin:
+  """Provides serialization methods for the Response dataclasses."""
+
+  def to_dict(self, exclude_none: bool = True) -> Dict[str, Any]:
+    """Converts the instance to a dictionary.
+
+        Args:
+            exclude_none: If True, only include non-empty values in the response.
+
+        Returns:
+            Dict[str, Any]: The dictionary representation of the instance.
+        """
+
+    def _remove_none(data: Any) -> Any:
+      """Recursively removes None or empty values from a dictionary or list."""
+      if isinstance(data, dict):
+        return {k: _remove_none(v) for k, v in data.items() if v is not None}
+      elif isinstance(data, list):
+        return [_remove_none(item) for item in data]
+      return data
+
+    result = asdict(self)
+    return _remove_none(result) if exclude_none else result
+
+  def to_json(self, exclude_none: bool = True) -> str:
+    """Converts the instance to a JSON string.
+
+        Args:
+            exclude_none: If True, only include non-empty values in the response.
+
+        Returns:
+            str: The JSON string representation of the instance.
+        """
+    return json.dumps(self.to_dict(exclude_none=exclude_none), indent=2)
 
 
 @dataclass
-class DCResponse:
+class DCResponse(SerializableMixin):
   """Represents a structured response from the Data Commons API."""
 
   json: Dict[str, Any] = field(default_factory=dict)
@@ -33,7 +69,7 @@ def next_token(self):
 
 
 @dataclass
-class NodeResponse:
+class NodeResponse(SerializableMixin):
   """Represents a response from the Node endpoint of the Data Commons API.
 
     Attributes:
@@ -69,13 +105,9 @@ def parse_data(data: Dict[str, Any]) -> Arcs | Properties:
   def get_properties(self) -> Dict:
     return flatten_properties(self.data)
 
-  @property
-  def json(self):
-    return asdict(self)
-
 
 @dataclass
-class ObservationResponse:
+class ObservationResponse(SerializableMixin):
   """Represents a response from the Observation endpoint of the Data Commons API.
 
     Attributes:
@@ -100,10 +132,6 @@ def from_json(cls, json_data: Dict[str, Any]) -> "ObservationResponse":
         },
     )
 
-  @property
-  def json(self):
-    return asdict(self)
-
   def get_data_by_entity(self) -> Dict:
     """Unpacks the data for each entity, for each variable.
 
@@ -125,7 +153,7 @@ def get_observations_as_records(self) -> List[Dict[str, Any]]:
 
 
 @dataclass
-class ResolveResponse:
+class ResolveResponse(SerializableMixin):
   """Represents a response from the Resolve endpoint of the Data Commons API.
 
     Attributes:
@@ -149,15 +177,3 @@ def from_json(cls, json_data: Dict[str, Any]) -> "ResolveResponse":
     return cls(entities=[
         Entity.from_json(entity) for entity in json_data.get("entities", [])
     ])
-
-  @property
-  def json(self):
-    """Converts the ResolveResponse instance to a dictionary.
-
-        This is useful for serializing the response data back into a JSON-compatible
-        format.
-
-        Returns:
-            Dict[str, Any]: The dictionary representation of the ResolveResponse instance.
-        """
-    return asdict(self)

datacommons_client/tests/endpoints/test_response.py

@@ -1,3 +1,5 @@
+import json
+
 from datacommons_client.endpoints.response import DCResponse
 from datacommons_client.endpoints.response import NodeResponse
 from datacommons_client.endpoints.response import ObservationResponse
@@ -78,7 +80,48 @@ def test_node_as_dict():
   }
 
   response = NodeResponse.from_json(json_data)
-  result = response.json
+  result = response.to_dict()
+
+  assert result == json_data
+
+
+def test_node_as_dict_exclude_none():
+  """Test that the NodeResponse.json property returns the correct dictionary."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "properties": None
+          }
+      },
+      "nextToken": "token123",
+  }
+
+  expected = {
+      "data": {
+          "geoId/06": {}
+      },
+      "nextToken": "token123",
+  }
+
+  response = NodeResponse.from_json(json_data)
+  result = response.to_dict(exclude_none=True)
+
+  assert result == expected
+
+
+def test_node_as_dict_include_none():
+  """Test that the NodeResponse.json property returns the correct dictionary."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "properties": None
+          }
+      },
+      "nextToken": "token123",
+  }
+
+  response = NodeResponse.from_json(json_data)
+  result = response.to_dict(exclude_none=False)
 
   assert result == json_data
 
@@ -238,14 +281,46 @@ def test_observation_as_dict():
   response = ObservationResponse.from_json(json_data)
 
   # Getting it back as a dictionary
-  result = response.json
+  result = response.to_dict()
 
   assert "byVariable" in result
   assert "facets" in result
   assert "var1" in result["byVariable"]
   assert "entity1" in result["byVariable"]["var1"]["byEntity"]
 
 
+def test_observation_as_dict_exclude_none():
+  """Test that the ObservationResponse.json property returns the correct dictionary."""
+  json_data = {
+      "byVariable": {
+          "var1": {
+              "byEntity": {
+                  "entity1": {
+                      "orderedFacets": [{
+                          "facetId": "facet1"
+                      }]
+                  }
+              }
+          }
+      },
+      "facets": {
+          "facet1": {
+              "unit": "GTQ",
+              "importName": "Import Name",
+          }
+      },
+  }
+
+  # Parsing JSON data
+  response = ObservationResponse.from_json(json_data)
+
+  # Getting it back as a dictionary
+  result = response.to_dict(exclude_none=True)
+
+  assert ("latestDate" not in result["byVariable"]["var1"]["byEntity"]
+          ["entity1"]["orderedFacets"][0])
+
+
 def test_observation_response_from_json():
   """Test that the ObservationResponse.from_json method parses JSON correctly."""
   json_data = {
@@ -483,8 +558,8 @@ def test_resolve_response_from_json():
   assert entity2.candidates[0].dominantType == "Type2"
 
 
-def test_resolve_response_json():
-  """Test that ResolveResponse.from_json and json are consistent."""
+def test_resolve_response_dict():
+  """Test that ResolveResponse.to_dict and json are consistent."""
   # Input dictionary
   input_data = {
       "entities": [
@@ -516,7 +591,118 @@ def test_resolve_response_json():
   response = ResolveResponse.from_json(input_data)
 
   # Convert back to dictionary using the json property
-  result = response.json
+  result = response.to_dict(exclude_none=False)
 
   # Assert that the resulting dictionary matches the original input
   assert result == input_data
+
+
+def test_resolve_response_dict_exclude_none():
+  """Test that ResolveResponse.to_dict and json are consistent."""
+  # Input dictionary
+  input_data = {
+      "entities": [{
+          "node":
+              "entity1",
+          "candidates": [
+              {
+                  "dcid": "dcid1",
+                  "dominantType": "Type1"
+              },
+              {
+                  "dcid": "dcid2",
+                  "dominantType": None
+              },
+          ],
+      }, {
+          "node": "entity2",
+          "candidates": [{
+              "dcid": "dcid3",
+              "dominantType": "Type2"
+          },],
+      }, {
+          "node": "entity3",
+          "candidates": [],
+      }]
+  }
+
+  # Expected data
+  expected_data = {
+      "entities": [{
+          "node":
+              "entity1",
+          "candidates": [
+              {
+                  "dcid": "dcid1",
+                  "dominantType": "Type1"
+              },
+              {
+                  "dcid": "dcid2"
+              },
+          ],
+      }, {
+          "node": "entity2",
+          "candidates": [{
+              "dcid": "dcid3",
+              "dominantType": "Type2"
+          },],
+      }, {
+          "node": "entity3",
+          "candidates": [],
+      }]
+  }
+
+  # Create ResolveResponse from the dictionary
+  response = ResolveResponse.from_json(input_data)
+
+  # Convert back to dictionary using the json property
+  result = response.to_dict(exclude_none=True)
+
+  # Assert that the resulting dictionary matches the original input
+  assert result == expected_data
+
+
+def test_resolve_response_json_string_exclude_none():
+  """Test that ResolveResponse.to_dict and json are consistent."""
+  # Input dictionary
+  input_data = {
+      "entities": [
+          {
+              "node":
+                  "entity1",
+              "candidates": [
+                  {
+                      "dcid": "dcid1",
+                      "dominantType": "Type1"
+                  },
+                  {
+                      "dcid": "dcid2",
+                      "dominantType": None
+                  },
+              ],
+          },
+          {
+              "node": "entity2",
+              "candidates": [{
+                  "dcid": "dcid3",
+                  "dominantType": "Type2"
+              },],
+          },
+          {
+              "node": "entity3",
+              "candidates": [],
+          },
+      ]
+  }
+
+  # Expected data
+  expected = json.dumps(input_data, indent=2)
+
+  # Create ResolveResponse from the dictionary
+  response = ResolveResponse.from_json(input_data)
+
+  # Convert back to dictionary using the json property
+  result = response.to_json(exclude_none=False)
+
+  # Assert that the resulting dictionary matches the original input
+  assert result == expected