Implement Data Classes to Model API Responses

datacommons_client/models/node.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, TypeAlias
+
+NextToken: TypeAlias = Optional[str]
+NodeDCID: TypeAlias = str
+ArcLabel: TypeAlias = str
+Property: TypeAlias = str
+PropertyList: TypeAlias = list[Property]
+
+
+@dataclass
+class Node:
+    """Represents an individual node in the Data Commons knowledge graph.
+
+    Attributes:
+        dcid: The unique identifier for the node.
+        name: The name of the node.
+        provenanceId: The provenance ID for the node.
+        types: The types associated with the node.
+        value: The value of the node.
+    """
+
+    dcid: str = None
+    name: Optional[str] = None
+    provenanceId: Optional[str] = None
+    types: Optional[list[str]] = None
+    value: Optional[str] = None
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Node":
+        return cls(
+            dcid=json_data.get("dcid"),
+            name=json_data.get("name"),
+            provenanceId=json_data.get("provenanceId"),
+            types=json_data.get("types"),
+            value=json_data.get("value"),
+        )
+
+
+@dataclass
+class NodeGroup:
+    """Represents a group of nodes in the Data Commons knowledge graph.
+
+    Attributes:
+        nodes: A list of Node objects in the group.
+    """
+
+    nodes: List[Node] = field(default_factory=list)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "NodeGroup":
+        """Parses a dictionary of lists of nodes from the response data.
+
+        Args:
+            json_data: The raw JSON data containing node information.
+
+        Returns:
+            A NodeGroup instance.
+        """
+        return cls(
+            nodes=[Node.from_json(node) for node in json_data.get("nodes", [])]
+        )
+
+
+@dataclass
+class Arcs:
+    """Represents arcs in the Data Commons knowledge graph.
+
+    Attributes:
+        arcs: A dictionary mapping arc labels to NodeGroup objects.
+    """
+
+    arcs: Dict[ArcLabel, NodeGroup] = field(default_factory=dict)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Arcs":
+        """Parses a dictionary of arcs from JSON.
+
+        Args:
+            json_data: The raw JSON data containing arc information.
+
+        Returns:
+            An Arcs instance.
+        """
+        return cls(
+            arcs={
+                label: NodeGroup.from_json(nodes)
+                for label, nodes in json_data.items()
+            }
+        )
+
+
+@dataclass
+class Properties:
+    """Represents a group of properties in the Data Commons knowledge graph.
+
+    Attributes:
+        properties: A list of property strings.
+    """
+
+    properties: List[Property] = field(default_factory=PropertyList)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Properties":
+        """Parses a list of properties from JSON.
+
+        Args:
+            json_data: The raw JSON data containing property information.
+
+        Returns:
+            A Properties instance.
+        """
+        return cls(properties=json_data.get("properties", []))

datacommons_client/models/observation.py

@@ -0,0 +1,147 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, TypeAlias
+
+variableDCID: TypeAlias = str
+facetID: TypeAlias = str
+
+orderedFacetsLabel: TypeAlias = str
+
+
+@dataclass
+class Observation:
+    """Represents an observation with a date and value.
+
+    Attributes:
+        date (str): The date of the observation.
+        value (float): The value of the observation.
+    """
+
+    date: str
+    value: float
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Observation":
+        """Creates an Observation instance from the response data.
+
+        Args:
+            json_data: A dictionary containing observation data.
+
+        Returns:
+            An Observation instance.
+        """
+        return cls(
+            date=json_data.get("date"),
+            value=json_data.get("value"),
+        )
+
+
+@dataclass
+class OrderedFacets:
+    """Represents ordered facets of observations.
+
+    Attributes:
+        earliestDate (str): The earliest date in the observations.
+        facetId (str): The identifier for the facet.
+        latestDate (str): The latest date in the observations.
+        obsCount (int): The total number of observations.
+        observations (List[Observation]): A list of observations associated with the facet.
+    """
+
+    earliestDate: str = field(default_factory=str)
+    facetId: str = field(default_factory=str)
+    latestDate: str = field(default_factory=str)
+    obsCount: int = field(default_factory=int)
+    observations: list[Observation] = field(default_factory=list)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "OrderedFacets":
+        """Creates an OrderedFacets instance from the response data.
+
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing ordered facets data.
+
+        Returns:
+            OrderedFacets: An instance of the OrderedFacets class.
+        """
+        return cls(
+            earliestDate=json_data.get("earliestDate"),
+            facetId=json_data.get("facetId"),
+            latestDate=json_data.get("latestDate"),
+            obsCount=json_data.get("obsCount"),
+            observations=[
+                Observation.from_json(observation)
+                for observation in json_data.get("observations", [])
+            ],
+        )
+
+
+@dataclass
+class Variable:
+    """Represents a variable with data grouped by entity.
+
+    Attributes:
+        byEntity (Dict[str, Dict[orderedFacetsLabel, List[OrderedFacets]]]): A dictionary mapping
+            entities to their ordered facets.
+    """
+
+    byEntity: Dict[str, Dict[orderedFacetsLabel, list[OrderedFacets]]] = field(
+        default_factory=dict
+    )
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Variable":
+        """Creates a Variable instance from the response data.
+
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing variable data.
+
+        Returns:
+            Variable: An instance of the Variable class.
+        """
+        return cls(
+            byEntity={
+                entity: {
+                    "orderedFacets": [
+                        OrderedFacets.from_json(facet_data)
+                        for facet_data in entity_data["orderedFacets"]
+                    ]
+                }
+                for entity, entity_data in json_data["byEntity"].items()
+            }
+        )
+
+
+@dataclass
+class Facet:
+    """Represents metadata for a facet.
+
+    Attributes:
+        importName (str): The name of the data import.
+        measurementMethod (str): The method used to measure the data.
+        observationPeriod (str): The period over which the observations were made.
+        provenanceUrl (str): The URL of the data's provenance.
+        unit (str): The unit of the observations.
+    """
+
+    importName: str = field(default_factory=str)
+    measurementMethod: str = field(default_factory=str)
+    observationPeriod: str = field(default_factory=str)
+    provenanceUrl: str = field(default_factory=str)
+    unit: str = field(default_factory=str)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Facet":
+        """
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing facet data.
+
+        Returns:
+            Facet: An instance of the Facet class.
+        """
+        return cls(
+            importName=json_data.get("importName"),
+            measurementMethod=json_data.get("measurementMethod"),
+            observationPeriod=json_data.get("observationPeriod"),
+            provenanceUrl=json_data.get("provenanceUrl"),
+            unit=json_data.get("unit"),
+        )

datacommons_client/models/resolve.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, TypeAlias
+
+Query: TypeAlias = str
+DCID: TypeAlias = str
+DominantType: TypeAlias = str
+
+
+@dataclass
+class Candidate:
+    """Represents a candidate in the resolution response.
+
+    Attributes:
+        dcid (DCID): The Data Commons ID for the candidate.
+        dominantType (Optional[DominantType]): The dominant type of the candidate,
+            if available. This represents the primary type associated with the DCID.
+    """
+
+    dcid: DCID = field(default_factory=str)
+    dominantType: Optional[DominantType] = None
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Candidate":
+        """Parses a Candidate instance from the response data.
+
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing candidate data,
+                typically from the Data Commons API.
+
+        Returns:
+            Candidate: An instance of the Candidate class populated with the
+            provided data.
+        """
+        return cls(
+            dcid=json_data["dcid"],
+            dominantType=json_data.get("dominantType"),
+        )
+
+
+@dataclass
+class Entity:
+    """Represents an entity with its resolution candidates.
+
+    Attributes:
+        node (Query): The query string or node being resolved.
+        candidates (List[Candidate]): A list of candidates that match the query.
+    """
+
+    node: Query
+    candidates: List[Candidate] = field(default_factory=list)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Entity":
+        """Parses an Entity instance from response data.
+
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing entity data,
+                including the node and associated candidates.
+
+        Returns:
+            Entity: A populated instance of the Entity class.
+        """
+        return cls(
+            node=json_data.get("node"),
+            candidates=[
+                Candidate.from_json(candidate)
+                for candidate in json_data.get("candidates", [])
+            ],
+        )

datacommons_client/models/sparql.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+
+@dataclass
+class Cell:
+    """Represents a single cell in a row.
+
+    Attributes:
+        value: The value contained in the cell.
+    """
+
+    value: str = field(default_factory=str)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Cell":
+        """Parses a Cell instance from response data.
+
+        Args:
+            json_data (Dict[str, Any]): A dictionary containing cell data.
+
+        Returns:
+            Cell: A populated instance of the Cell class.
+        """
+
+        return cls(value=json_data.get("value"))
+
+
+@dataclass
+class Row:
+    """Parses a Row instance from response data.
+
+    Args:
+        json_data (Dict[str, Any]): A dictionary containing row data,
+            typically with a "cells" key containing a list of cells.
+
+    Returns:
+        Row: A populated instance of the Row class.
+    """
+
+    cells: List[Cell] = field(default_factory=list)
+
+    @classmethod
+    def from_json(cls, json_data: Dict[str, Any]) -> "Row":
+        return cls(
+            cells=[Cell.from_json(cell) for cell in json_data.get("cells", [])]
+        )