refactor types

Author: jm-rivera

datacommons_client/endpoints/node.py

@@ -6,13 +6,12 @@
 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.models.graph import Parent
 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
 
@@ -197,85 +196,95 @@ def fetch_all_classes(
     )
 
   def fetch_entity_parents(
-      self, entity_dcids: str | list[str]) -> dict[str, list[Parent]]:
+      self,
+      entity_dcids: str | list[str],
+      *,
+      as_dict: bool = True) -> dict[str, list[Parent | dict]]:
     """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.
+            as_dict (bool): If True, returns a dictionary mapping each input DCID to its
+                immediate parent entities. If False, returns a dictionary of Parent objects (which
+                are dataclasses).
 
         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.
+            dict[str, list[Parent | dict]]: 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) or as a dictionary with
+            the same data.
         """
     # 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)
+    result = build_parents_dictionary(data=data)
 
+    if as_dict:
+      return {k: [p.to_dict() for p in v] for k, v in result.items()}
 
-def _fetch_parents_cached(self, dcid: str) -> tuple[Parent, ...]:
-  """Returns cached parent nodes for a given entity using an LRU cache.
+    return result
 
-    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.
+  def _fetch_parents_cached(self, dcid: str) -> tuple[Parent, ...]:
+    """Returns cached parent nodes for a given entity using an LRU cache.
 
-    Args:
-        dcid (str): The DCID of the entity whose parents should be fetched.
+        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.
 
-    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'.
-    """
+        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
+    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

datacommons_client/models/graph.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+from typing import TypeAlias
+
+from datacommons_client.utils.data_processing import SerializableMixin
+
+
+@dataclass(frozen=True)
+class Parent(SerializableMixin):
+  """A class representing a parent node in a graph.
+    Attributes:
+        dcid (str): The ID of the parent node.
+        name (str): The name of the parent node.
+        type (str | list[str]): The type(s) of the parent node.
+    """
+
+  dcid: str
+  name: str
+  type: str | list[str]
+
+
+# A dictionary mapping DCIDs to lists of Parent objects.
+AncestryMap: TypeAlias = dict[str, list[Parent]]

datacommons_client/utils/graph.py

@@ -3,31 +3,14 @@
 from concurrent.futures import Future
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import wait
-from dataclasses import dataclass
 from functools import lru_cache
 from typing import Callable, Optional
 
-from datacommons_client.utils.data_processing import SerializableMixin
+from datacommons_client.models.graph import AncestryMap
+from datacommons_client.models.graph import Parent
 
 PARENTS_MAX_WORKERS = 10
 
-
-@dataclass(frozen=True)
-class Parent(SerializableMixin):
-  """A class representing a parent node in a graph.
-    Attributes:
-        dcid (str): The ID of the parent node.
-        name (str): The name of the parent node.
-        type (str | list[str]): The type(s) of the parent node.
-    """
-
-  dcid: str
-  name: str
-  type: str | list[str]
-
-
-AncestryMap = dict[str, list[Parent]]
-
 # -- -- Fetch tools -- --
 
 
@@ -48,7 +31,7 @@ def _fetch_parents_uncached(endpoint, dcid: str) -> list[Parent]:
     Returns:
         A list of parent dictionaries, each containing 'dcid', 'name', and 'type'.
     """
-  return endpoint.fetch_entity_parents(dcid).get(dcid, [])
+  return endpoint.fetch_entity_parents(dcid, as_dict=False).get(dcid, [])
 
 
 @lru_cache(maxsize=512)