Update node.py

Add ancestry tools

Author: jm-rivera

datacommons_client/endpoints/node.py

@@ -1,10 +1,20 @@
+from concurrent.futures import ThreadPoolExecutor
 from typing import Optional
 
 from datacommons_client.endpoints.base import API
 from datacommons_client.endpoints.base import Endpoint
 from datacommons_client.endpoints.payloads import NodeRequestPayload
 from datacommons_client.endpoints.payloads import normalize_properties_to_string
 from datacommons_client.endpoints.response import NodeResponse
+from datacommons_client.models.node import Node
+from datacommons_client.utils.graph import build_ancestry_map
+from datacommons_client.utils.graph import build_ancestry_tree
+from datacommons_client.utils.graph import build_parents_dictionary
+from datacommons_client.utils.graph import fetch_parents_lru
+from datacommons_client.utils.graph import flatten_ancestry
+from datacommons_client.utils.graph import Parent
+
+ANCESTRY_MAX_WORKERS = 20
 
 
 class NodeEndpoint(Endpoint):
@@ -91,10 +101,12 @@ def fetch_property_labels(
     expression = "->" if out else "<-"
 
     # Make the request and return the response.
-    return self.fetch(node_dcids=node_dcids,
-                      expression=expression,
-                      all_pages=all_pages,
-                      next_token=next_token)
+    return self.fetch(
+        node_dcids=node_dcids,
+        expression=expression,
+        all_pages=all_pages,
+        next_token=next_token,
+    )
 
   def fetch_property_values(
       self,
@@ -143,10 +155,12 @@ def fetch_property_values(
     if constraints:
       expression += f"{{{constraints}}}"
 
-    return self.fetch(node_dcids=node_dcids,
-                      expression=expression,
-                      all_pages=all_pages,
-                      next_token=next_token)
+    return self.fetch(
+        node_dcids=node_dcids,
+        expression=expression,
+        all_pages=all_pages,
+        next_token=next_token,
+    )
 
   def fetch_all_classes(
       self,
@@ -174,8 +188,94 @@ def fetch_all_classes(
             ```
         """
 
-    return self.fetch_property_values(node_dcids="Class",
-                                      properties="typeOf",
-                                      out=False,
-                                      all_pages=all_pages,
-                                      next_token=next_token)
+    return self.fetch_property_values(
+        node_dcids="Class",
+        properties="typeOf",
+        out=False,
+        all_pages=all_pages,
+        next_token=next_token,
+    )
+
+  def fetch_entity_parents(
+      self, entity_dcids: str | list[str]) -> dict[str, list[Parent]]:
+    """Fetches the direct parents of one or more entities using the 'containedInPlace' property.
+
+        Args:
+            entity_dcids (str | list[str]): A single DCID or a list of DCIDs to query.
+
+        Returns:
+            dict[str, list[Parent]]: A dictionary mapping each input DCID to a list of its
+            immediate parent entities. Each parent is represented as a Parent object, which
+            contains the DCID, name, and type of the parent entity.
+        """
+    # Fetch property values from the API
+    data = self.fetch_property_values(
+        node_dcids=entity_dcids,
+        properties="containedInPlace",
+    ).get_properties()
+
+    return build_parents_dictionary(data=data)
+
+
+def _fetch_parents_cached(self, dcid: str) -> tuple[Parent, ...]:
+  """Returns cached parent nodes for a given entity using an LRU cache.
+
+    This private wrapper exists because `@lru_cache` cannot be applied directly
+    to instance methods. By passing the `NodeEndpoint` instance (`self`) as an
+    argument caching is preserved while keeping the implementation modular and testable.
+
+    Args:
+        dcid (str): The DCID of the entity whose parents should be fetched.
+
+    Returns:
+        tuple[Parent, ...]: A tuple of Parent objects representing the entity's immediate parents.
+    """
+  return fetch_parents_lru(self, dcid)
+
+
+def fetch_entity_ancestry(
+    self,
+    entity_dcids: str | list[str],
+    as_tree: bool = False) -> dict[str, list[dict[str, str]] | dict]:
+  """Fetches the full ancestry (flat or nested) for one or more entities.
+    For each input DCID, this method builds the complete ancestry graph using a
+    breadth-first traversal and parallel fetching.
+    It returns either a flat list of unique parents or a nested tree structure for
+    each entity, depending on the `as_tree` flag. The flat list matches the structure
+    of the `/api/place/parent` endpoint of the DC website.
+    Args:
+        entity_dcids (str | list[str]): One or more DCIDs of the entities whose ancestry
+           will be fetched.
+        as_tree (bool): If True, returns a nested tree structure; otherwise, returns a flat list.
+            Defaults to False.
+    Returns:
+        dict[str, list[dict[str, str]] | dict]: A dictionary mapping each input DCID to either:
+            - A flat list of parent dictionaries (if `as_tree` is False), or
+            - A nested ancestry tree (if `as_tree` is True). Each parent is represented by
+              a dict with 'dcid', 'name', and 'type'.
+    """
+
+  if isinstance(entity_dcids, str):
+    entity_dcids = [entity_dcids]
+
+  result = {}
+
+  # Use a thread pool to fetch ancestry graphs in parallel for each input entity
+  with ThreadPoolExecutor(max_workers=ANCESTRY_MAX_WORKERS) as executor:
+    futures = [
+        executor.submit(build_ancestry_map,
+                        root=dcid,
+                        fetch_fn=self._fetch_parents_cached)
+        for dcid in entity_dcids
+    ]
+
+    # Gather ancestry maps and postprocess into flat or nested form
+    for future in futures:
+      dcid, ancestry = future.result()
+      if as_tree:
+        ancestry = build_ancestry_tree(dcid, ancestry)
+      else:
+        ancestry = flatten_ancestry(ancestry)
+      result[dcid] = ancestry
+
+  return result