Add EndpointPayload ABC and Specific Payload Dataclasses for API Requests (#208)

jm-rivera · web-flow · commit 0c1f1a5ad137 · 2025-01-17T11:45:52.000-06:00
This PR adds 'payload' dataclasses for the API endpoints.

The EndpointPayload class is an Abstract Base Class, which defines the structure for all other payload classes. Payload classes are in charge of normalizing and validating the payload before it is used for a POST request. They also implement a .to_dict property for easy access to the correctly structured payload data as a python dictionary.
diff --git a/datacommons_client/endpoints/payloads.py b/datacommons_client/endpoints/payloads.py
@@ -0,0 +1,218 @@
+from abc import ABC
+from abc import abstractmethod
+from dataclasses import dataclass
+from dataclasses import field
+from enum import Enum
+from typing import Optional
+
+from datacommons_client.utils.error_handling import (
+    InvalidObservationSelectError,
+)
+
+
+@dataclass
+class EndpointRequestPayload(ABC):
+    """
+    Abstract base class for payload dataclasses.
+    Defines the required interface for all payload dataclasses.
+    """
+
+    @abstractmethod
+    def normalize(self) -> None:
+        """Normalize the payload for consistent internal representation."""
+        pass
+
+    @abstractmethod
+    def validate(self) -> None:
+        """Validate the payload to ensure its structure and contents are correct."""
+        pass
+
+    @property
+    @abstractmethod
+    def to_dict(self) -> dict:
+        """Convert the payload into a dictionary format for API requests."""
+        pass
+
+
+@dataclass
+class NodeRequestPayload(EndpointRequestPayload):
+    """
+    A dataclass to structure, normalize, and validate the payload for a Node V2 API request.
+
+    Attributes:
+        node_dcids (str | list[str]): The DCID(s) of the nodes to query.
+        expression (str): The property or relation expression(s) to query.
+    """
+
+    node_dcids: str | list[str]
+    expression: str
+
+    def __post_init__(self):
+        self.normalize()
+        self.validate()
+
+    def normalize(self):
+        if isinstance(self.node_dcids, str):
+            self.node_dcids = [self.node_dcids]
+
+    def validate(self):
+        if not isinstance(self.expression, str):
+            raise ValueError("Expression must be a string.")
+
+    @property
+    def to_dict(self) -> dict:
+        return {"nodes": self.node_dcids, "property": self.expression}
+
+
+class ObservationSelect(str, Enum):
+    DATE = "date"
+    VARIABLE = "variable"
+    ENTITY = "entity"
+    VALUE = "value"
+
+    @classmethod
+    def _missing_(cls, value):
+        """Handle missing enum values by raising a custom error."""
+        valid_values = [member.value for member in cls]
+        message = f"Invalid `select` field: '{value}'. Only {', '.join(valid_values)} are allowed."
+        raise InvalidObservationSelectError(message=message)
+
+
+class ObservationDate(str, Enum):
+    LATEST = "LATEST"
+    ALL = ""
+
+
+@dataclass
+class ObservationRequestPayload(EndpointRequestPayload):
+    """
+    A dataclass to structure, normalize, and validate the payload for an Observation V2 API request.
+
+    Attributes:
+        date (str): The date for which data is being requested.
+        variable_dcids (str | list[str]): One or more variable IDs for the data.
+        select (list[ObservationSelect]): Fields to include in the response.
+            Defaults to ["date", "variable", "entity", "value"].
+        entity_dcids (Optional[str | list[str]]): One or more entity IDs to filter the data.
+        entity_expression (Optional[str]): A string expression to filter entities.
+    """
+
+    date: ObservationDate | str = ""
+    variable_dcids: str | list[str] = field(default_factory=list)
+    select: list[ObservationSelect | str] = field(
+        default_factory=lambda: [
+            ObservationSelect.DATE,
+            ObservationSelect.VARIABLE,
+            ObservationSelect.ENTITY,
+            ObservationSelect.VALUE,
+        ]
+    )
+    entity_dcids: Optional[str | list[str]] = None
+    entity_expression: Optional[str] = None
+
+    def __post_init__(self):
+        """
+        Initializes the payload, performing validation and normalization.
+
+        Raises:
+            ValueError: If validation rules are violated.
+        """
+        self.RequiredSelect = {"variable", "entity"}
+        self.normalize()
+        self.validate()
+
+    def normalize(self):
+        """
+        Normalizes the payload for consistent internal representation.
+
+        - Converts `variable_dcids` and `entity_dcids` to lists if they are passed as strings.
+        - Normalizes the `date` field to ensure it is in the correct format.
+        """
+        # Normalize variable
+        if isinstance(self.variable_dcids, str):
+            self.variable_dcids = [self.variable_dcids]
+
+        # Normalize entity
+        if isinstance(self.entity_dcids, str):
+            self.entity_dcids = [self.entity_dcids]
+
+        # Normalize date field
+        if self.date.upper() == "ALL":
+            self.date = ObservationDate.ALL
+        elif (self.date.upper() == "LATEST") or (self.date == ""):
+            self.date = ObservationDate.LATEST
+
+    def validate(self):
+        """
+        Validates the payload to ensure consistency and correctness.
+
+        Raises:
+            ValueError: If both `entity_dcids` and `entity_expression` are set,
+                        if neither is set, or if required fields are missing from `select`.
+        """
+
+        # Validate mutually exclusive entity fields
+        if bool(self.entity_dcids) == bool(self.entity_expression):
+            raise ValueError(
+                "Exactly one of 'entity_dcids' or 'entity_expression' must be set."
+            )
+
+        # Check if all required fields are present
+        missing_fields = self.RequiredSelect - set(self.select)
+        if missing_fields:
+            raise ValueError(
+                f"The 'select' field must include at least the following: {', '.join(self.RequiredSelect)} "
+                f"(missing: {', '.join(missing_fields)})"
+            )
+
+        # Check all select fields are valid
+        [ObservationSelect(select_field) for select_field in self.select]
+
+    @property
+    def to_dict(self) -> dict:
+        """
+        Converts the payload into a dictionary format for API requests.
+
+        Returns:
+            dict: The normalized and validated payload.
+        """
+        return {
+            "date": self.date,
+            "variable": {"dcids": self.variable_dcids},
+            "entity": (
+                {"dcids": self.entity_dcids}
+                if self.entity_dcids
+                else {"expression": self.entity_expression}
+            ),
+            "select": self.select,
+        }
+
+
+@dataclass
+class ResolveRequestPayload(EndpointRequestPayload):
+    """
+    A dataclass to structure, normalize, and validate the payload for a Resolve V2 API request.
+
+    Attributes:
+        node_dcids (str | list[str]): The DCID(s) of the nodes to query.
+        expression (str): The relation expression to query.
+    """
+
+    node_dcids: str | list[str]
+    expression: str
+
+    def __post_init__(self):
+        self.normalize()
+        self.validate()
+
+    def normalize(self):
+        if isinstance(self.node_dcids, str):
+            self.node_dcids = [self.node_dcids]
+
+    def validate(self):
+        if not isinstance(self.expression, str):
+            raise ValueError("Expression must be a string.")
+
+    @property
+    def to_dict(self) -> dict:
+        return {"nodes": self.node_dcids, "property": self.expression}
diff --git a/datacommons_client/tests/endpoints/test_payloads.py b/datacommons_client/tests/endpoints/test_payloads.py
@@ -0,0 +1,134 @@
+from datacommons_client.endpoints.payloads import NodeRequestPayload
+from datacommons_client.endpoints.payloads import ObservationDate
+from datacommons_client.endpoints.payloads import ObservationRequestPayload
+from datacommons_client.endpoints.payloads import ObservationSelect
+from datacommons_client.endpoints.payloads import ResolveRequestPayload
+from datacommons_client.utils.error_handling import (
+    InvalidObservationSelectError,
+)
+import pytest
+
+
+def test_node_payload_normalize():
+    """Tests that NodeRequestPayload correctly normalizes single and multiple node_dcids."""
+    payload = NodeRequestPayload(node_dcids="node1", expression="prop1")
+    assert payload.node_dcids == ["node1"]
+
+    payload = NodeRequestPayload(node_dcids=["node1", "node2"], expression="prop1")
+    assert payload.node_dcids == ["node1", "node2"]
+
+
+def test_node_payload_validate():
+    """Tests that NodeRequestPayload validates its inputs correctly."""
+    with pytest.raises(ValueError):
+        NodeRequestPayload(
+            node_dcids="node1", expression=123
+        )  # `expression` must be a string
+
+
+def test_node_payload_to_dict():
+    """Tests NodeRequestPayload conversion to dictionary."""
+    payload = NodeRequestPayload(node_dcids="node1", expression="prop1")
+    assert payload.to_dict == {"nodes": ["node1"], "property": "prop1"}
+
+
+def test_observation_payload_normalize():
+    """Tests that ObservationRequestPayload normalizes inputs correctly."""
+    payload = ObservationRequestPayload(
+        date="LATEST",
+        variable_dcids="var1",
+        select=["variable", "entity"],
+        entity_dcids="ent1",
+    )
+    assert payload.variable_dcids == ["var1"]
+    assert payload.entity_dcids == ["ent1"]
+    assert payload.date == ObservationDate.LATEST
+
+    payload = ObservationRequestPayload(
+        date="all",
+        variable_dcids=["var1"],
+        select=["variable", "entity"],
+        entity_dcids=["ent1"],
+    )
+    assert payload.date == ObservationDate.ALL
+    assert payload.variable_dcids == ["var1"]
+    assert payload.entity_dcids == ["ent1"]
+
+
+def test_observation_select_invalid_value():
+    """Tests that an invalid ObservationSelect value raises InvalidObservationSelectError."""
+    with pytest.raises(
+        InvalidObservationSelectError,
+        match=r"Invalid `select` field: 'invalid'. Only date, variable, entity, value are allowed.",
+    ):
+        ObservationSelect("invalid")
+
+
+def test_observation_payload_validate():
+    """Tests that ObservationRequestPayload validates its inputs."""
+    with pytest.raises(ValueError):
+        ObservationRequestPayload(
+            date="LATEST",
+            variable_dcids="var1",
+            select=["variable"],
+            entity_dcids=None,
+            entity_expression=None,
+        )  # Requires either `entity_dcids` or `entity_expression`
+
+    with pytest.raises(ValueError):
+        ObservationRequestPayload(
+            date="LATEST",
+            variable_dcids="var1",
+            select=["value"],  # Missing required "variable" and "entity"
+            entity_expression="expression",
+        )
+
+    with pytest.raises(ValueError):
+        ObservationRequestPayload(
+            date="LATEST",
+            variable_dcids="var1",
+            select=["variable", "entity"],
+            entity_dcids="ent1",
+            entity_expression="expression",  # Both `entity_dcids` and `entity_expression` set
+        )
+
+
+def test_observation_payload_to_dict():
+    """Tests ObservationRequestPayload conversion to dictionary."""
+    payload = ObservationRequestPayload(
+        date="LATEST",
+        variable_dcids="var1",
+        select=["variable", "entity"],
+        entity_dcids="ent1",
+    )
+    assert payload.to_dict == {
+        "date": ObservationDate.LATEST,
+        "variable": {"dcids": ["var1"]},
+        "entity": {"dcids": ["ent1"]},
+        "select": ["variable", "entity"],
+    }
+
+
+def test_resolve_payload_normalize():
+    """Tests that ResolveRequestPayload normalizes single and multiple node_dcids."""
+    payload = ResolveRequestPayload(node_dcids="node1", expression="expr1")
+    assert payload.node_dcids == ["node1"]
+
+    payload = ResolveRequestPayload(
+        node_dcids=["node1", "node2"], expression="expr1"
+    )
+    assert payload.node_dcids == ["node1", "node2"]
+
+
+def test_resolve_payload_validate():
+    """Tests that ResolveRequestPayload validates its inputs correctly."""
+    with pytest.raises(ValueError):
+        ResolveRequestPayload(
+            node_dcids="node1", expression=123
+        )  # `expression` must be a string
+
+
+def test_resolve_payload_to_dict():
+    """Tests ResolveRequestPayload conversion to dictionary."""
+    payload = ResolveRequestPayload(node_dcids="node1", expression="expr1")
+    assert payload.to_dict == {"nodes": ["node1"], "property": "expr1"}
diff --git a/datacommons_client/tests/endpoints/test_request_handling.py b/datacommons_client/tests/endpoints/test_request_handling.py
@@ -4,12 +4,10 @@
 import pytest
 import requests
 
-from datacommons_client.utils.error_hanlding import (
-    DCAuthenticationError,
-    DCConnectionError,
-    DCStatusError,
-    APIError,
-)
+from datacommons_client.utils.error_hanlding import APIError
+from datacommons_client.utils.error_hanlding import DCAuthenticationError
+from datacommons_client.utils.error_hanlding import DCConnectionError
+from datacommons_client.utils.error_hanlding import DCStatusError
 from datacommons_client.utils.error_hanlding import InvalidDCInstanceError
 from datacommons_client.utils.request_handling import _check_instance_is_valid
 from datacommons_client.utils.request_handling import _fetch_with_pagination
diff --git a/datacommons_client/utils/error_handling.py b/datacommons_client/utils/error_handling.py
@@ -76,3 +76,9 @@ class InvalidDCInstanceError(DataCommonsError):
     """Raised when an invalid Data Commons instance is provided."""
 
     default_message = "The specified Data Commons instance is invalid."
+
+
+class InvalidObservationSelectError(DataCommonsError):
+    """Raised when an invalid ObservationSelect field is provided."""
+
+    default_message = "The ObservationSelect field is invalid."
