Add client class (#214)

Introduces the`DataCommonsClient` class to interact with the Data Commons API, along with various utility functions and tests.

Author: jm-rivera

datacommons_client/client.py

@@ -0,0 +1,106 @@
+from typing import Literal, Optional
+
+from datacommons_client.endpoints.base import API
+from datacommons_client.endpoints.node import NodeEndpoint
+from datacommons_client.endpoints.observation import ObservationEndpoint
+from datacommons_client.endpoints.payloads import ObservationDate
+from datacommons_client.endpoints.resolve import ResolveEndpoint
+from datacommons_client.utils.decorators import requires_pandas
+
+try:
+  import pandas as pd
+except ImportError:
+  pd = None
+
+
+class DataCommonsClient:
+  """
+    A client for interacting with the Data Commons API.
+
+    This class provides convenient access to the V2 Data Commons API endpoints.
+
+    Attributes:
+        api (API): An instance of the API class that handles requests.
+        node (NodeEndpoint): Provides access to node-related queries, such as fetching property labels
+            and values for individual or multiple nodes in the Data Commons knowledge graph.
+        observation (ObservationEndpoint): Handles observation-related queries, allowing retrieval of
+            statistical observations associated with entities, variables, and dates (e.g., GDP of California in 2010).
+        resolve (ResolveEndpoint): Manages resolution queries to find different DCIDs for entities.
+
+    """
+
+  def __init__(
+      self,
+      api_key: Optional[str] = None,
+      *,
+      dc_instance: Optional[str] = "datacommons.org",
+      url: Optional[str] = None,
+  ):
+    """
+        Initializes the DataCommonsClient.
+
+        Args:
+            api_key (Optional[str]): The API key for authentication. Defaults to None. Note that
+                custom DC instances do not currently require an API key.
+            dc_instance (Optional[str]): The Data Commons instance to use. Defaults to "datacommons.org".
+            url (Optional[str]): A custom, fully resolved URL for the Data Commons API. Defaults to None.
+        """
+    # Create an instance of the API class which will be injected to the endpoints
+    self.api = API(api_key=api_key, dc_instance=dc_instance, url=url)
+
+    # Create instances of the endpoints
+    self.node = NodeEndpoint(api=self.api)
+    self.observation = ObservationEndpoint(api=self.api)
+    self.resolve = ResolveEndpoint(api=self.api)
+
+  @requires_pandas
+  def observations_dataframe(
+      self,
+      variable_dcids: str | list[str],
+      date: ObservationDate | str,
+      entity_dcids: Literal["all"] | list[str] = "all",
+      entity_type: Optional[str] = None,
+      parent_entity: Optional[str] = None,
+  ):
+    """
+        Fetches statistical observations and returns them as a Pandas DataFrame.
+
+        The Observation API fetches statistical observations linked to entities and variables
+        at a particular date (e.g., "population of USA in 2020", "GDP of California in 2010").
+
+        Args:
+            variable_dcids (str | list[str]): One or more variable DCIDs for the observation.
+            date (ObservationDate | str): The date for which observations are requested. It can be
+            a specific date, "all" to retrieve all observations, or "latest" to get the most recent observations.
+            entity_dcids (Literal["all"] | list[str], optional): The entity DCIDs to retrieve data for.
+                Defaults to "all". DCIDs must include their type (e.g "country/GTM" for Guatemala).
+            entity_type (Optional[str], optional): The type of entities to filter by when `entity_dcids="all"`.
+                Required if `entity_dcids="all"`. Defaults to None.
+            parent_entity (Optional[str], optional): The parent entity under which the target entities fall.
+                Used only when `entity_dcids="all"`. Defaults to None.
+
+        Returns:
+            pd.DataFrame: A DataFrame containing the requested observations.
+        """
+
+    if entity_dcids == "all" and not entity_type:
+      raise ValueError(
+          "When 'entity_dcids' is 'all', 'entity_type' must be specified.")
+
+    if entity_dcids != "all" and (entity_type or parent_entity):
+      raise ValueError(
+          "Specify 'entity_type' and 'parent_entity' only when 'entity_dcids' is 'all'."
+      )
+
+    if entity_dcids == "all":
+      observations = self.observation.fetch_observations_by_entity_type(
+          variable_dcids=variable_dcids,
+          date=date,
+          entity_type=entity_type,
+          parent_entity=parent_entity,
+      )
+    else:
+      observations = self.observation.fetch_observations_by_entity(
+          variable_dcids=variable_dcids, date=date, entity_dcids=entity_dcids)
+
+    return pd.DataFrame(observations.get_observations_as_records())

datacommons_client/endpoints/observation.py

@@ -104,7 +104,7 @@ def fetch_observations_by_entity_type(
       parent_entity: str,
       entity_type: str,
       variable_dcids: str | list[str],
-  ):
+  ) -> ObservationResponse:
     """
         Fetches all observations for a given entity type.
 
@@ -143,3 +143,43 @@ def fetch_observations_by_entity_type(
         entity_expression=
         f"{parent_entity}<-containedInPlace+{{typeOf:{entity_type}}}",
     )
+
+  def fetch_observations_by_entity(
+      self,
+      date: ObservationDate | str,
+      entity_dcids: str | list[str],
+      variable_dcids: str | list[str],
+  ) -> ObservationResponse:
+    """
+        Fetches all observations for a given entity type.
+
+        Args:
+            date (ObservationDate | str): The date option for the observations.
+                Use 'all' for all dates, 'latest' for the most recent data,
+                or provide a date as a string (e.g., "2024").
+            entity_dcids (str | list[str]): One or more entity IDs to filter the data.
+            variable_dcids (str | list[str]): The variable(s) to fetch observations for.
+                This can be a single variable ID or a list of IDs.
+
+        Returns:
+            ObservationResponse: The response object containing observations for the specified entity type.
+
+        Example:
+            To fetch all observations for Nigeria for a specific variable:
+
+            ```python
+            api = API()
+            ObservationEndpoint(api).fetch_observations_by_entity(
+                date="all",
+                entity_dcids="country/NGA",
+                variable_dcids="sdg/SI_POV_DAY1"
+            )
+            ```
+        """
+
+    return self.fetch(
+        variable_dcids=variable_dcids,
+        date=date,
+        entity_dcids=entity_dcids,
+        select=[s for s in ObservationSelect],
+    )

datacommons_client/endpoints/response.py

@@ -12,6 +12,9 @@
 from datacommons_client.models.observation import Variable
 from datacommons_client.models.observation import variableDCID
 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
 
 
 @dataclass
@@ -71,55 +74,6 @@ def json(self):
     return asdict(self)
 
 
-def flatten_properties(data: Dict[str, Any]) -> Dict[str, Any]:
-  """
-    Flatten the properties of a node response.
-
-    Processes a dictionary of node responses, extracting and
-    simplifying their properties and arcs into a flattened dictionary.
-
-    Args:
-        data (Dict[str, Dict[str, Any]]):
-            The input dictionary containing node responses. Each node maps to
-            a dictionary with potential "arcs" and "properties" keys.
-
-    Returns:
-        Dict[str, Any]:
-            A flattened dictionary where keys are node identifiers, and values
-            are the simplified properties or nodes.
-    """
-
-  # Store simplified properties
-  items = {}
-
-  for node, node_data in data.items():
-    # If arcs are present, process them
-    if hasattr(node_data, "arcs"):
-      processed_arcs = _unpack_arcs(node_data.arcs)
-      if processed_arcs is not None:
-        items[node] = processed_arcs
-        continue
-
-    # Include properties if present
-    if hasattr(node_data, "properties"):
-      items[node] = node_data.properties
-
-  return items
-
-
-def _unpack_arcs(arcs: Dict[str, Any]) -> Any:
-  """Simplify the 'arcs' structure."""
-  if len(arcs) > 1:
-    # Multiple arcs: return dictionary of property nodes
-    return {prop: arc_data["nodes"] for prop, arc_data in arcs.items()}
-
-  # Single arc: extract first node's data
-  for property_data in arcs.values():
-    nodes = property_data.nodes
-    if nodes is not None:
-      return nodes if len(nodes) > 1 else nodes[0]
-
-
 @dataclass
 class ObservationResponse:
   """Represents a response from the Observation endpoint of the Data Commons API.
@@ -170,61 +124,6 @@ def get_observations_as_records(self) -> List[Dict[str, Any]]:
                                    facets=self.facets)
 
 
-def extract_observations(variable: str, entity: str, entity_data: dict,
-                         facet_metadata: dict) -> list[dict]:
-  """
-    Extracts observations for a given variable, entity, and its data.
-
-    Args:
-        variable (str): The variable name.
-        entity (str): The entity name.
-        entity_data (dict): Data for the entity, including ordered facets.
-        facet_metadata (dict): Metadata for facets.
-
-    Returns:
-        list[dict]: A list of observation records.
-    """
-  # Store observation records
-  records = []
-
-  # Extract observations
-  for facet in entity_data.get("orderedFacets", []):
-    facet_id = facet.facetId
-    metadata = facet_metadata.get(facet_id, {})
-    records.extend({
-        "date": observation.date,
-        "entity": entity,
-        "variable": variable,
-        "value": observation.value,
-        "facetId": facet_id,
-        **asdict(metadata),
-    } for observation in facet.observations)
-  return records
-
-
-def observations_as_records(data: dict, facets: dict) -> list[dict]:
-  """
-    Converts observation data into a list of records.
-
-    Args:
-        data (dict): A mapping of variables to entities and their data.
-        facets (dict): Facet metadata for the observations.
-
-    Returns:
-        list[dict]: A flattened list of observation records.
-    """
-  return [
-      record for variable, entities in data.items()
-      for entity, entity_data in entities.items()
-      for record in extract_observations(
-          variable=variable,
-          entity=entity,
-          entity_data=entity_data,
-          facet_metadata=facets,
-      )
-  ]
-
-
 @dataclass
 class ResolveResponse:
   """Represents a response from the Resolve endpoint of the Data Commons API.

datacommons_client/models/observation.py

@@ -105,7 +105,7 @@ def from_json(cls, json_data: Dict[str, Any]) -> "Variable":
                     OrderedFacets.from_json(facet_data)
                     for facet_data in entity_data.get("orderedFacets", {})
                 ]
-            } for entity, entity_data in json_data["byEntity"].items()
+            } for entity, entity_data in json_data.get("byEntity", {}).items()
         })
 
 

datacommons_client/tests/endpoints/test_response.py

@@ -1,14 +1,14 @@
-from datacommons_client.endpoints.response import _unpack_arcs
 from datacommons_client.endpoints.response import DCResponse
-from datacommons_client.endpoints.response import extract_observations
-from datacommons_client.endpoints.response import flatten_properties
 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.observation import Facet
 from datacommons_client.models.observation import Observation
 from datacommons_client.models.observation import OrderedFacets
 from datacommons_client.models.observation import Variable
+from datacommons_client.utils.data_processing import extract_observations
+from datacommons_client.utils.data_processing import flatten_properties
+from datacommons_client.utils.data_processing import unpack_arcs
 
 ### ----- Test DCResponse ----- ###
 
@@ -153,7 +153,7 @@ def test_unpack_arcs_multiple_properties():
       },  # Empty nodes for completeness
   }
 
-  result = _unpack_arcs(arcs)
+  result = unpack_arcs(arcs)
 
   # Expected output
   expected = {