Add fetch_available_statistical_variables (#229)

**Fetch available statistical variables**: 
- Adds a `fetch_available_statistical_variables` method to the
ObservationEndpoint. This method fetches all statvars which have
observations for one or more entities.

Author: jm-rivera

datacommons_client/endpoints/observation.py

@@ -6,6 +6,7 @@
 from datacommons_client.endpoints.payloads import ObservationRequestPayload
 from datacommons_client.endpoints.payloads import ObservationSelect
 from datacommons_client.endpoints.response import ObservationResponse
+from datacommons_client.utils.data_processing import group_variables_by_entity
 
 
 class ObservationEndpoint(Endpoint):
@@ -233,3 +234,25 @@ def fetch_observations_by_entity(
         filter_facet_domains=filter_facet_domains,
         filter_facet_ids=filter_facet_ids,
     )
+
+  def fetch_available_statistical_variables(
+      self,
+      entity_dcids: str | list[str],
+  ) -> dict[str, list[str]]:
+    """
+        Fetches available statistical variables (which have observations) for given entities.
+        Args:
+            entity_dcids (str | list[str]): One or more entity DCIDs(s) to fetch variables for.
+        Returns:
+            dict[str, list[str]]: A dictionary mapping entity DCIDs to their available statistical variables.
+        """
+
+    # Fetch observations for the given entity DCIDs. If variable is empty list
+    # all available variables are retrieved.
+    data = self.fetch(
+        entity_dcids=entity_dcids,
+        select=[ObservationSelect.VARIABLE, ObservationSelect.ENTITY],
+        variable_dcids=[],
+    ).get_data_by_entity()
+
+    return group_variables_by_entity(data=data)

datacommons_client/tests/endpoints/test_observation_endpoint.py

@@ -3,6 +3,7 @@
 from datacommons_client.endpoints.base import API
 from datacommons_client.endpoints.observation import ObservationEndpoint
 from datacommons_client.endpoints.payloads import ObservationDate
+from datacommons_client.endpoints.payloads import ObservationSelect
 from datacommons_client.endpoints.response import ObservationResponse
 
 
@@ -164,3 +165,51 @@ def test_fetch_observations_facets_by_entity_type():
                                         endpoint="observation",
                                         all_pages=True,
                                         next_token=None)
+
+
+def test_fetch_available_statistical_variables_single_entity():
+  """Test fetching variables for a single entity."""
+  mock_data = {
+      "var1": ["ent1"],
+      "var2": ["ent1"],
+  }
+
+  # Mock the fetch method on the ObservationEndpoint instance
+  endpoint = ObservationEndpoint(api=MagicMock())
+  endpoint.fetch = MagicMock()
+  endpoint.fetch.return_value.get_data_by_entity = MagicMock(
+      return_value=mock_data)
+
+  result = endpoint.fetch_available_statistical_variables("ent1")
+
+  expected = {
+      "ent1": ["var1", "var2"],
+  }
+  assert result == expected
+
+  endpoint.fetch.assert_called_once_with(
+      entity_dcids="ent1",
+      select=[ObservationSelect.VARIABLE, ObservationSelect.ENTITY],
+      variable_dcids=[],
+  )
+
+
+def test_fetch_available_statistical_variables_multiple_entities():
+  """Test fetching variables for multiple entities."""
+  mock_data = {
+      "var1": ["ent1", "ent2"],
+      "var2": ["ent2"],
+  }
+
+  endpoint = ObservationEndpoint(api=MagicMock())
+  endpoint.fetch = MagicMock()
+  endpoint.fetch.return_value.get_data_by_entity = MagicMock(
+      return_value=mock_data)
+
+  result = endpoint.fetch_available_statistical_variables(["ent1", "ent2"])
+
+  expected = {
+      "ent1": ["var1"],
+      "ent2": ["var1", "var2"],
+  }
+  assert result == expected

datacommons_client/tests/test_utils.py

@@ -0,0 +1,40 @@
+from datacommons_client.utils.data_processing import group_variables_by_entity
+
+
+def test_group_variables_by_entity_basic():
+  """Test grouping with simple variable-entity mapping."""
+  input_data = {
+      "var1": ["ent1", "ent2"],
+      "var2": ["ent2", "ent3"],
+      "var3": ["ent1"],
+  }
+  expected_output = {
+      "ent1": ["var1", "var3"],
+      "ent2": ["var1", "var2"],
+      "ent3": ["var2"],
+  }
+
+  result = group_variables_by_entity(input_data)
+  assert result == expected_output
+
+
+def test_group_variables_by_entity_duplicate_entities():
+  """Test grouping when a variable has duplicate entities."""
+  input_data = {
+      "var1": ["ent1", "ent1", "ent2"],
+  }
+  result = group_variables_by_entity(input_data)
+  assert result["ent1"].count("var1") == 2  # duplicates are preserved
+  assert "ent2" in result
+  assert result["ent2"] == ["var1"]
+
+
+def test_group_variables_by_entity_preserves_order():
+  """Test if the order of variables is preserved in the resulting entity lists."""
+  input_data = {
+      "var1": ["ent1"],
+      "var2": ["ent1"],
+      "var3": ["ent1"],
+  }
+  result = group_variables_by_entity(input_data)
+  assert result["ent1"] == ["var1", "var2", "var3"]

datacommons_client/utils/data_processing.py

@@ -93,3 +93,23 @@ def observations_as_records(data: dict, facets: dict) -> list[dict]:
           facet_metadata=facets,
       )
   ]
+
+
+def group_variables_by_entity(
+    data: dict[str, list[str]]) -> dict[str, list[str]]:
+  """Groups variables by the entities they are associated with.
+      Takes a dictionary mapping statistical variable DCIDs to a list of entity DCIDs,
+      and returns a new dictionary mapping each entity DCID to a list of statistical
+      variables available for that entity.
+      Args:
+          data: A dictionary where each key is a variable DCID and the value is a list
+              of entity DCIDs that have observations for that variable.
+      Returns:
+          A dictionary where each key is an entity DCID and the value is a list of
+          variable DCIDs available for that entity.
+      """
+  result: dict[str, list[str]] = {}
+  for variable, entities in data.items():
+    for entity in entities:
+      result.setdefault(entity, []).append(variable)
+  return result