Add helpers to extract data from NodeResponse arcs (#246)

This adds two helper methods to simplify reading data from a
NodeResponse.

`extract_connected_dcids`:
* An example use case: Starting from a gene, I want to see associated
diseases and indicators for that accociations. To accomplish this, I
must find all of the dcids for the gene-disease-association nodes, then
perform a follow up query for properties of those dcids. Using this
helper method, it would look like:
    ```
    node_client = dc_client.node
    resp = node_client.fetch("bio/APOE", "<-geneID")
gene_association_dcids = resp.extract_connected_dcids("bio/APOE",
"geneID", connected_node_types="DiseaseGeneAssociation")
association_metdata = node_client.fetch(gene_association_dcids,
"->[relationshipAssociationType,relationshipEvidenceType, confidence]")
    ```
* The helper method makes it easy to extract just the dcids from the
response that I care about without having to really dive into what the
structure of the response is.

`extract_connected_nodes`: This has a similar use case to the first
helper method, but returns the nodes for the arc that I'm interested in,
instead of just the DCIDs.

---
This is a revised version of closed
#238

@jm-rivera Let me know what you think! I can provide some more use cases
if that would be helpful.

Would it be more helpful to rename these "extract_nodes_from_arc" /
"extract_connected_dcids_from_arc"? Granted, those are pretty long
names....

Author: clincoln8

datacommons_client/endpoints/response.py

@@ -1,12 +1,15 @@
 from dataclasses import dataclass
 from dataclasses import field
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from datacommons_client.models.base import SerializableMixin
 from datacommons_client.models.node import Arcs
 from datacommons_client.models.node import NextToken
+from datacommons_client.models.node import Node
 from datacommons_client.models.node import NodeDCID
+from datacommons_client.models.node import NodeGroup
 from datacommons_client.models.node import Properties
+from datacommons_client.models.node import Property
 from datacommons_client.models.observation import Facet
 from datacommons_client.models.observation import facetID
 from datacommons_client.models.observation import Variable
@@ -68,6 +71,72 @@ def parse_data(data: Dict[str, Any]) -> Arcs | Properties:
   def get_properties(self) -> Dict:
     return flatten_properties(self.data)
 
+  def extract_connected_nodes(
+      self,
+      subject_dcid: NodeDCID,
+      property_dcid: Property,
+      connected_node_types: Optional[str | list[str]] = None) -> List[Node]:
+    """Retrieves Node objects in the NodeResponse connected to the subject node
+    via the specified property.
+
+    Args:
+      subject_dcid: The DCID of the starting node in the arc.
+      property_dcid: The property connecting the subject node to the desired
+        target nodes.
+      connected_node_types: Optional. A type or list of types to filter the
+        connected nodes. If provided, only connected nodes that have at least
+        one of the specified types will be returned. If omitted, all nodes from
+        the arc are returned.
+
+    Returns:
+      A list of Node objects that are connected to the subject node via the
+      specified property.
+    """
+    if isinstance(connected_node_types, str):
+      connected_node_types = [connected_node_types]
+
+    nodes = self.get_properties().get(subject_dcid, {}).get(property_dcid, [])
+
+    connected_nodes = []
+    for node in nodes:
+      if connected_node_types:
+        # Filter out nodes that are missing a list of types or do not have the
+        # desired type
+        if not node.types or not any(nt in node.types
+                                     for nt in connected_node_types):
+          continue
+
+      connected_nodes.append(node)
+
+    return connected_nodes
+
+  def extract_connected_dcids(
+      self,
+      subject_dcid: NodeDCID,
+      property_dcid: Property,
+      connected_node_types: Optional[str | list[str]] = None) -> List[NodeDCID]:
+    """Retrieves DCIDs of the Nodes in the NodeResponse connected to the subject
+    node via the specified property.
+
+    Args:
+      subject_dcid: The DCID of the starting node.
+      property_dcid: The property connecting the subject node to the desired
+        target nodes.
+      connected_node_types: Optional. A type or list of types to filter the
+        connected nodes. If provided, only DCIDs of connected nodes that have at
+        least one of the specified types will be returned. If omitted, DCIDs of
+        all nodes from the arc are returned.
+
+    Returns:
+      A list of NodeDCIDs for the nodes connected via the specified property
+      from the subject node.
+    """
+
+    connected_nodes = self.extract_connected_nodes(subject_dcid, property_dcid,
+                                                   connected_node_types)
+
+    return [node.dcid for node in connected_nodes if node.dcid]
+
 
 @dataclass
 class ObservationResponse(SerializableMixin):

datacommons_client/tests/endpoints/test_response.py

@@ -5,6 +5,7 @@
 from datacommons_client.endpoints.response import NodeResponse
 from datacommons_client.endpoints.response import ObservationResponse
 from datacommons_client.endpoints.response import ResolveResponse
+from datacommons_client.models.node import Arcs
 from datacommons_client.models.node import Node
 from datacommons_client.models.node import NodeGroup
 from datacommons_client.models.observation import Facet
@@ -287,6 +288,234 @@ def test_unpack_arcs_multiple_properties():
   assert result == expected
 
 
+def test_extract_connected_dcids():
+  """Test that extract_connected_dcids is successful when multiple dcid and multiple
+  properties are in the response."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "containedInPlace": {
+                      "nodes": [{
+                          "dcid": "country/USA",
+                          "name": "United States",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["Country"]
+                      }, {
+                          "dcid": "usc/PacificDivision",
+                          "name": "Pacific Division",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["CensusDivision"]
+                      }]
+                  },
+                  "name": {
+                      "nodes": [{
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "value": "California"
+                      }]
+                  }
+              }
+          },
+          "geoId/07": {
+              "arcs": {
+                  "containedInPlace": {
+                      "nodes": [{
+                          "dcid": "country/USA",
+                          "name": "United States",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["Country"]
+                      }]
+                  },
+              }
+          }
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(subject_dcid='geoId/06',
+                                            property_dcid='containedInPlace')
+  assert result == ['country/USA', 'usc/PacificDivision']
+
+
+def test_extract_connected_dcids_with_nonexistent_dcid():
+  """Test that extract_connected_dcids returns empty when requested dcid is not in the
+  NodeResponse."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "name": {
+                      "nodes": [{
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "value": "California"
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(subject_dcid='geoId/07',
+                                            property_dcid='name')
+  assert result == []
+
+
+def test_extract_connected_dcids_with_nonexistent_property():
+  """Test that extract_connected_dcids returns empty when requested property is not in
+  the NodeResponse."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "name": {
+                      "nodes": [{
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "value": "California"
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(subject_dcid='geoId/06',
+                                            property_dcid='containedInPlace')
+  assert result == []
+
+
+def test_extract_connected_dcids_does_not_include_none_for_value_only_nodes():
+  """Test that extract_connected_dcids does not include None in the returned list
+  when the nodes in the response only contain values and not dcids."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "name": {
+                      "nodes": [{
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "value": "California"
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(subject_dcid='geoId/06',
+                                            property_dcid='name')
+  assert result == []
+
+
+def test_extract_connected_dcids_with_node_type_filter():
+  """Test that extract_connected_dcids returns dcids with the corresponding
+  node_type."""
+
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "relatedPlaces": {
+                      "nodes": [{
+                          "dcid": "country/USA",
+                          "name": "United States",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["Country"]
+                      }, {
+                          "dcid": "usc/PacificDivision",
+                          "name": "Pacific Division",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["CensusDivision"]
+                      }, {
+                          "dcid": "node3",
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(subject_dcid='geoId/06',
+                                            property_dcid='relatedPlaces',
+                                            connected_node_types="Country")
+  assert result == ['country/USA']
+
+
+def test_extract_connected_dcids_with_multiple_node_type_filter():
+  """Test that extract_connected_dcids returns dcids with the corresponding
+  connected_node_types."""
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "relatedPlaces": {
+                      "nodes": [{
+                          "dcid": "country/USA",
+                          "name": "United States",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["Country"]
+                      }, {
+                          "dcid": "usc/PacificDivision",
+                          "name": "Pacific Division",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["CensusDivision"]
+                      }, {
+                          "dcid": "node3",
+                          "types": ["City"]
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_dcids(
+      subject_dcid='geoId/06',
+      property_dcid='relatedPlaces',
+      connected_node_types=["Country", "City"])
+  assert result == ['country/USA', 'node3']
+
+
+def test_extract_connected_nodes_with_multiple_node_type_filter():
+  """Test that extract_connected_nodes returns only nodes with the corresponding
+  connected_node_types."""
+
+  json_data = {
+      "data": {
+          "geoId/06": {
+              "arcs": {
+                  "relatedPlaces": {
+                      "nodes": [{
+                          "dcid": "country/USA",
+                          "name": "United States",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["Country"]
+                      }, {
+                          "dcid": "usc/PacificDivision",
+                          "name": "Pacific Division",
+                          "provenanceId": "dc/base/WikidataOtherIdGeos",
+                          "types": ["CensusDivision"]
+                      }, {
+                          "dcid": "node3",
+                          "types": ["City"]
+                      }]
+                  }
+              }
+          },
+      }
+  }
+  response = NodeResponse.from_json(json_data)
+  result = response.extract_connected_nodes(
+      subject_dcid='geoId/06',
+      property_dcid='relatedPlaces',
+      connected_node_types=["Country", "City"])
+  assert result == [
+      Node(dcid="country/USA",
+           name="United States",
+           provenanceId="dc/base/WikidataOtherIdGeos",
+           types=["Country"]),
+      Node(dcid="node3", types=["City"])
+  ]
+
+
 ### ----- Test Observation Response ----- ###