Optionally include constraint property metadata in observation DataFrames (#274)

This PR adds support for including property constraint metadata for
statistical variables in observation DataFrames. To support this, it
includes new models and endpoint methods to fetch and represent these
constraints. The PR also includes tests for all new functionality.

In short:
* Added new pydantic models `StatVarConstraint` and `StatVarConstraints`
to represent constraints associated with statistical variables.
* Implemented `NodeEndpoint.fetch_statvar_constraints` to fetch and
structure constraint property/value pairs for a set of statvars.
* Updated `Client.observations_dataframe` to support an
`include_constraints_metadata` argument; when enabled, the returned
DataFrame includes constraint id/name columns for each variable.
* Added utility function
`add_property_constraints_to_observations_dataframe` to insert
constraint columns into observation DataFrames.

Author: jm-rivera

datacommons_client/client.py

@@ -6,6 +6,7 @@
 from datacommons_client.endpoints.resolve import ResolveEndpoint
 from datacommons_client.models.observation import ObservationDate
 from datacommons_client.utils.dataframes import add_entity_names_to_observations_dataframe
+from datacommons_client.utils.dataframes import add_property_constraints_to_observations_dataframe
 from datacommons_client.utils.decorators import requires_pandas
 from datacommons_client.utils.error_handling import NoDataForPropertyError
 
@@ -92,7 +93,8 @@ def _find_filter_facet_ids(
           date=date,
           entity_dcids=entity_dcids,
           variable_dcids=variable_dcids,
-          select=["variable", "entity", "facet"])
+          select=["variable", "entity", "facet"],
+      )
     else:
       observations = self.observation.fetch_observations_by_entity_type(
           date=date,
@@ -120,6 +122,7 @@ def observations_dataframe(
       entity_type: Optional[str] = None,
       parent_entity: Optional[str] = None,
       property_filters: Optional[dict[str, str | list[str]]] = None,
+      include_constraints_metadata: bool = False,
   ):
     """
         Fetches statistical observations and returns them as a Pandas DataFrame.
@@ -139,6 +142,9 @@ def observations_dataframe(
             Required if `entity_dcids="all"`. Defaults to None.
         property_filters (Optional[dict[str, str | list[str]]): An optional dictionary used to filter
             the data by using observation properties like `measurementMethod`, `unit`, or `observationPeriod`.
+        include_constraints_metadata (bool): If True, includes the dcid and name of any constraint
+            properties associated with the variable DCIDs (based on the `constraintProperties` property)
+            in the returned DataFrame. Defaults to False.
 
         Returns:
             pd.DataFrame: A DataFrame containing the requested observations.
@@ -181,7 +187,8 @@ def observations_dataframe(
           date=date,
           entity_dcids=entity_dcids,
           variable_dcids=variable_dcids,
-          filter_facet_ids=facets)
+          filter_facet_ids=facets,
+      )
 
     # Convert the observations to a DataFrame
     df = pd.DataFrame(observations.to_observation_records().model_dump())
@@ -193,4 +200,10 @@ def observations_dataframe(
         entity_columns=["entity", "variable"],
     )
 
+    if include_constraints_metadata:
+      df = add_property_constraints_to_observations_dataframe(
+          endpoint=self.node,
+          observations_df=df,
+      )
+
     return df

datacommons_client/endpoints/node.py

@@ -11,6 +11,8 @@
 from datacommons_client.endpoints.response import NodeResponse
 from datacommons_client.models.node import Name
 from datacommons_client.models.node import Node
+from datacommons_client.models.node import StatVarConstraint
+from datacommons_client.models.node import StatVarConstraints
 from datacommons_client.utils.graph import build_graph_map
 from datacommons_client.utils.graph import build_relationship_tree
 from datacommons_client.utils.graph import fetch_relationship_lru
@@ -23,6 +25,8 @@
 
 PLACES_MAX_WORKERS = 10
 
+CONSTRAINT_PROPERTY: str = "constraintProperties"
+
 _DEPRECATED_METHODS: dict[str, dict[str, str | dict[str, str]]] = {
     "fetch_entity_parents": {
         "new_name": "fetch_place_parents",
@@ -534,3 +538,113 @@ def fetch_place_descendants(
         relationship="children",
         max_concurrent_requests=max_concurrent_requests,
     )
+
+  def _fetch_property_id_names(self, node_dcids: str | list[str],
+                               properties: str | list[str]):
+    """Fetch target nodes for given properties and return only (dcid, name).
+
+        For each input node and each requested property, returns the list of target
+        nodes as dictionaries with ``dcid`` and ``name``.
+
+        Args:
+            node_dcids: A single DCID or a list of DCIDs to query.
+            properties: A property string or list of property strings.
+
+        Returns:
+            A mapping:
+            `{ node_dcid: { property: [ {dcid, name}, ... ], ... }, ... }`.
+        """
+    data = self.fetch_property_values(node_dcids=node_dcids,
+                                      properties=properties).get_properties()
+
+    result: dict[str, dict[str, list[dict]]] = {}
+
+    for node, props in data.items():
+      result.setdefault(node, {})
+      for prop, metadata in props.items():
+        dest = result[node].setdefault(prop, [])
+        for n in metadata:
+          # Prefer 'dcid', but if property is terminal, fall back to 'value'.
+          dcid = n.dcid or n.value
+          name = n.name or n.value
+          dest.append({"dcid": dcid, "name": name})
+    return result
+
+  def fetch_statvar_constraints(
+      self, variable_dcids: str | list[str]) -> StatVarConstraints:
+    """Fetch constraint property/value pairs for statistical variables, using
+        the `constraintProperties` property.
+
+        This returns, for each StatisticalVariable, the constraints that define it.
+
+        Args:
+            variable_dcids: One or more StatisticalVariable DCIDs.
+
+        Returns:
+            StatVarConstraints:
+            ``{
+                <sv_dcid>: [
+                    {
+                        "constraint_id": <constraint_property_dcid>,
+                        "constraint_name": <constraint_property_name>,
+                        "value_id": <value_node_dcid>,
+                        "value_name": <value_node_name>,
+                    },
+                    ...
+                ],
+                ...
+            }``
+        """
+    # Ensure variable_dcids is a list
+    if isinstance(variable_dcids, str):
+      variable_dcids = [variable_dcids]
+
+    # Get constraints for the given variable DCIDs.
+    constraints_mapping = self._fetch_property_id_names(
+        node_dcids=variable_dcids, properties=[CONSTRAINT_PROPERTY])
+
+    # Per statvar mapping of dcid - name
+    per_sv_constraint_names = {}
+    # Global set of all constraint property IDs
+    all_constraint_prop_ids = set()
+
+    for sv in variable_dcids:
+      # Get the constraint properties for this statvar
+      prop_entries = constraints_mapping.get(sv,
+                                             {}).get(CONSTRAINT_PROPERTY, [])
+      # Map the constraint properties to their names
+      id_to_name = {entry["dcid"]: entry.get("name") for entry in prop_entries}
+      # Add an entry for this statvar to the constraint names mapping
+      per_sv_constraint_names[sv] = id_to_name
+      # Update the global set of all constraint property IDs
+      all_constraint_prop_ids.update(id_to_name.keys())
+
+    # In a single request, fetch all values for all the constraints, for all statvars.
+    values_map = self._fetch_property_id_names(
+        node_dcids=variable_dcids,
+        properties=sorted(all_constraint_prop_ids),
+    )
+
+    # Build structured response. This will include vars with no constraints (empty dicts).
+    result = {sv: [] for sv in variable_dcids}
+
+    for sv in variable_dcids:
+      constraint_names = per_sv_constraint_names.get(sv, {})
+      sv_values = values_map.get(sv, {})
+
+      for constraintId, constraintName in constraint_names.items():
+        values = sv_values.get(constraintId, [])
+        # Continue if the stat var doesn't actually define a value for one of its constraintProperties.
+        if not values:
+          continue
+
+        # Build the StatVarConstraint object
+        result[sv].append(
+            StatVarConstraint(
+                constraintId=constraintId,
+                constraintName=constraintName,
+                valueId=values[0]["dcid"],
+                valueName=values[0].get("name"),
+            ))
+
+    return StatVarConstraints.model_validate(result)

datacommons_client/models/node.py

@@ -90,3 +90,20 @@ class NodeList(BaseDCModel, ListLikeRootModel[list[Node]]):
 
 class NodeDCIDList(BaseDCModel, ListLikeRootModel[list[NodeDCID]]):
   """A root model whose value is a list of NodeDCID strings."""
+
+
+class StatVarConstraint(BaseDCModel):
+  """Represents a constraint for a statistical variable."""
+
+  constraintId: NodeDCID
+  constraintName: Optional[str] = None
+  valueId: NodeDCID
+  valueName: Optional[str] = None
+
+
+class StatVarConstraints(BaseDCModel,
+                         DictLikeRootModel[dict[NodeDCID,
+                                                list[StatVarConstraint]]]):
+  """A root model whose value is a dictionary of statvar ids - a list of StatVarConstraint objects.
+    This model is used to represent constraints associated with statistical variables.
+    """

datacommons_client/tests/endpoints/test_node_endpoint.py

@@ -8,6 +8,7 @@
 from datacommons_client.models.node import Name
 from datacommons_client.models.node import Node
 from datacommons_client.models.node import NodeGroup
+from datacommons_client.models.node import StatVarConstraints
 from datacommons_client.utils.names import DEFAULT_NAME_PROPERTY
 from datacommons_client.utils.names import NAME_WITH_LANGUAGE_PROPERTY
 
@@ -395,3 +396,171 @@ def test_fetch_entity_ancestry_tree(mock_build_map, mock_build_tree):
   mock_build_tree.assert_called_once_with(root="Y",
                                           graph=mock_build_map.return_value[1],
                                           relationship_key="parents")
+
+
+def test__fetch_property_id_names_flattens_to_dcid_and_name():
+  """Private helper should return only dcid and name per target node."""
+  api_mock = MagicMock(spec=API)
+  endpoint = NodeEndpoint(api=api_mock)
+
+  # Simulate fetch_property_values response with Arcs, NodeGroup, Node
+  endpoint.fetch_property_values = MagicMock(return_value=NodeResponse(
+      data={
+          "sv/1":
+              Arcs(
+                  arcs={
+                      "constraintProperties":
+                          NodeGroup(nodes=[
+                              Node(dcid="p1", name="Prop One"),
+                              Node(dcid="p2", name="Prop Two"),
+                          ])
+                  })
+      }))
+
+  result = endpoint._fetch_property_id_names("sv/1", "constraintProperties")
+
+  assert result == {
+      "sv/1": {
+          "constraintProperties": [
+              {
+                  "dcid": "p1",
+                  "name": "Prop One",
+              },
+              {
+                  "dcid": "p2",
+                  "name": "Prop Two",
+              },
+          ]
+      }
+  }
+  endpoint.fetch_property_values.assert_called_once_with(
+      node_dcids="sv/1", properties="constraintProperties")
+
+
+def test_fetch_statvar_constraints_builds_constraints_and_values():
+  """fetch_statvar_constraints should combine constraint properties and values."""
+  endpoint = NodeEndpoint(api=MagicMock())
+
+  # First call returns constraint property ids and names
+  constraints_map = {
+      "sv/1": {
+          "constraintProperties": [
+              {
+                  "dcid": "p1",
+                  "name": "Prop One"
+              },
+              {
+                  "dcid": "p2",
+                  "name": "Prop Two"
+              },
+          ]
+      }
+  }
+
+  # Second call returns values for those properties
+  values_map = {
+      "sv/1": {
+          "p1": [{
+              "dcid": "v1",
+              "name": "Val One"
+          }],
+          "p2": [{
+              "dcid": "v2",
+              "name": "Val Two"
+          }],
+      }
+  }
+
+  with patch.object(endpoint,
+                    "_fetch_property_id_names",
+                    side_effect=[constraints_map, values_map]) as mock_helper:
+    result = endpoint.fetch_statvar_constraints(["sv/1"])
+
+  # Ensure helper called twice (once for constraintProperties, once for values)
+  assert mock_helper.call_count == 2
+  assert isinstance(result, StatVarConstraints)
+  assert "sv/1" in result
+  # Two constraints returned
+  assert len(result["sv/1"]) == 2
+  ids = {(c.constraintId, c.valueId) for c in result["sv/1"]}
+  assert ids == {("p1", "v1"), ("p2", "v2")}
+
+
+def test_fetch_statvar_constraints_handles_string_input_and_no_constraints():
+  """Single sv input and empty constraints should yield empty list for that sv."""
+  endpoint = NodeEndpoint(api=MagicMock())
+
+  # No constraintProperties for sv/empty
+  constraints_map = {"sv/empty": {"constraintProperties": []}}
+  # Second call won't be used but provide empty map
+  values_map = {"sv/empty": {}}
+
+  with patch.object(endpoint,
+                    "_fetch_property_id_names",
+                    side_effect=[constraints_map, values_map]):
+    # string input
+    result = endpoint.fetch_statvar_constraints("sv/empty")
+
+  assert isinstance(result, StatVarConstraints)
+  assert result["sv/empty"] == []
+
+
+def test__fetch_property_id_names_handles_literal_values():
+  """_fetch_property_id_names should handle string literal values gracefully."""
+  api_mock = MagicMock(spec=API)
+  endpoint = NodeEndpoint(api=api_mock)
+
+  # Simulate a response where the target value is a literal string (no dcid)
+  endpoint.fetch_property_values = MagicMock(return_value=NodeResponse(
+      data={
+          "sv/1":
+              Arcs(arcs={"p1": NodeGroup(nodes=[Node(value="LiteralValue")])})
+      }))
+
+  result = endpoint._fetch_property_id_names("sv/1", "p1")
+
+  assert result == {
+      "sv/1": {
+          "p1": [{
+              "dcid": "LiteralValue",
+              "name": "LiteralValue"
+          }]
+      }
+  }
+  endpoint.fetch_property_values.assert_called_once_with(node_dcids="sv/1",
+                                                         properties="p1")
+
+
+def test_fetch_statvar_constraints_skips_missing_constraint_values():
+  """If a constraintProperty has no value for a SV, skip it without error."""
+  endpoint = NodeEndpoint(api=MagicMock())
+
+  constraints_map = {
+      "sv/1": {
+          "constraintProperties": [
+              {
+                  "dcid": "p1",
+                  "name": "Prop One"
+              },
+              {
+                  "dcid": "p2",
+                  "name": "Prop Two"
+              },
+          ]
+      }
+  }
+
+  # p1 has a value, p2 is missing/empty
+  values_map = {"sv/1": {"p1": [{"dcid": "v1", "name": "Val One"}], "p2": []}}
+
+  with patch.object(endpoint,
+                    "_fetch_property_id_names",
+                    side_effect=[constraints_map, values_map]):
+    result = endpoint.fetch_statvar_constraints(["sv/1"])
+
+  assert isinstance(result, StatVarConstraints)
+  assert "sv/1" in result
+  # Only one well-formed constraint should be included (p1)
+  assert len(result["sv/1"]) == 1
+  assert result["sv/1"][0].constraintId == "p1"
+  assert result["sv/1"][0].valueId == "v1"