Use pydantic to validate input data

pyproject.toml

@@ -1,6 +1,27 @@
+[project]
+name = "nerc_rates"
+authors = [
+{name="MOC Alliance"},
+]
+version = "0.1"
+readme = "README.md"
+classifiers = [
+    "Programming Language :: Python :: 3"
+]
+requires-python = ">=3.8"
+dependencies = [
+  "pydantic",
+  "pyyaml",
+  "requests",
+]
+
 [build-system]
 requires = [
-    "setuptools>=42",
-    "wheel"
+  "setuptools>=42",
+  "wheel"
 ]
 build-backend = "setuptools.build_meta"
+
+[tool.setuptools.package-dir]
+"" = "src"
+

rates.yaml

@@ -1,56 +1,50 @@
 #################################################
 # Service Unit Rates
 #################################################
-- name: CPU SU Rate
-  history:
-    - value: 0.013
-      from: 2023-06
-
-- name: GPUA100 SU Rate
-  history:
-    - value: 1.803
-      from: 2023-06
-
-- name: GPUA100SXM4 SU Rate
-  history:
-    - value: 2.078
-      from: 2023-06
-
-- name: GPUV100 SU Rate
-  history:
-    - value: 1.214
-      from: 2023-06
-
-- name: GPUK80 SU Rate
-  history:
-    - value: 0.463
-      from: 2023-06
-
-- name: Storage GB Rate
-  history:
-    - value: 0.000009
-      from: 2023-06
+- history:
+    - from: 2023-06
+      value: '0.013'
+  name: CPU SU Rate
+- history:
+    - from: 2023-06
+      value: '1.803'
+  name: GPUA100 SU Rate
+- history:
+    - from: 2023-06
+      value: '2.078'
+  name: GPUA100SXM4 SU Rate
+- history:
+    - from: 2023-06
+      value: '1.214'
+  name: GPUV100 SU Rate
+- history:
+    - from: 2023-06
+      value: '0.463'
+  name: GPUK80 SU Rate
+- history:
+    - from: 2023-06
+      value: '0.000009'
+  name: Storage GB Rate
 
 #################################################
 # Feature Flags
 #################################################
-- name: Charge for Stopped Instances
-  history:
-    - value: False
-      from: 2023-06
+- history:
+    - from: 2023-06
       until: 2024-02
-    - value: True
-      from: 2024-03
+      value: 'False'
+    - from: 2024-03
+      value: 'True'
+  name: Charge for Stopped Instances
 
 #################################################
 # SU Definitions
 #################################################
-- name: vCPUs in CPU SU
-  history:
-    - value: 1
-      from: 2023-06
-
-- name: RAM in CPU SU
-  history:
-    - value: 4096
-      from: 2023-06
+- history:
+    - from: 2023-06
+      value: '1'
+  name: vCPUs in CPU SU
+- history:
+    - from: 2023-06
+      value: '4096'
+  name: RAM in CPU SU

requirements.txt

@@ -1,2 +1,3 @@
+pydantic
 pyyaml
 requests

src/nerc_rates/models.py

@@ -0,0 +1,50 @@
+from typing import Annotated
+
+import datetime
+import pydantic
+
+
+def parse_date(v: str | datetime.date) -> datetime.date:
+    if isinstance(v, str):
+        return datetime.datetime.strptime(v, "%Y-%m").date()
+    return v
+
+
+DateField = Annotated[datetime.date, pydantic.BeforeValidator(parse_date)]
+
+
+class Base(pydantic.BaseModel):
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+
+class RateValue(Base):
+    value: str
+    date_from: Annotated[DateField, pydantic.Field(alias="from")]
+    date_until: Annotated[DateField, pydantic.Field(alias="until", default=None)]
+
+
+class RateItem(Base):
+    name: str
+    history: list[RateValue]
+
+
+RateItemDict = Annotated[
+    dict[str, RateItem],
+    pydantic.BeforeValidator(lambda items: {x["name"]: x for x in items}),
+]
+
+
+class Rates(pydantic.RootModel):
+    root: RateItemDict
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+    def get_value_at(self, name: str, queried_date: datetime.date | str):
+        d = parse_date(queried_date)
+        for item in self[name]["history"]:
+            if item.date_from <= d <= (item.date_until or d):
+                return item["value"]
+
+        raise ValueError(f"No value for {name} for {queried_date}.")

src/nerc_rates/rates.py

@@ -1,43 +1,22 @@
-from datetime import date, datetime
-
 import requests
 import yaml
 
-DEFAULT_URL = "https://raw.githubusercontent.com/knikolla/nerc-rates/main/rates.yaml"
-
-
-class Rates:
-    def __init__(self, config):
-        self.values = {x["name"]: x for x in config}
+from .models import Rates
 
-    @staticmethod
-    def _parse_date(d: str | date) -> date:
-        if isinstance(d, str):
-            d = datetime.strptime(d, "%Y-%m").date()
-        return d
-
-    def get_value_at(self, name: str, queried_date: date | str):
-        d = self._parse_date(queried_date)
-        for v_dict in self.values[name]["history"]:
-            v_from = self._parse_date(v_dict["from"])
-            v_until = self._parse_date(v_dict.get("until", d))
-            if v_from <= d <= v_until:
-                return v_dict["value"]
-
-        raise ValueError(f"No value for {name} for {queried_date}.")
+DEFAULT_URL = "https://raw.githubusercontent.com/knikolla/nerc-rates/main/rates.yaml"
 
 
 def load_from_url(url=DEFAULT_URL) -> Rates:
     r = requests.get(url, allow_redirects=True)
     # Using the BaseLoader prevents conversion of numeric
     # values to floats and loads them as strings.
-    config = yaml.load(r.content.decode("utf-8"), Loader=yaml.BaseLoader)
-    return Rates(config)
+    config = yaml.safe_load(r.content.decode("utf-8"))
+    return Rates.model_validate(config)
 
 
 def load_from_file() -> Rates:
     with open("rates.yaml", "r") as f:
         # Using the BaseLoader prevents conversion of numeric
         # values to floats and loads them as strings.
-        config = yaml.load(f, Loader=yaml.BaseLoader)
-    return Rates(config)
+        config = yaml.safe_load(f)
+    return Rates.model_validate(config)

src/nerc_rates/tests/test_rates.py

@@ -8,7 +8,7 @@ def test_load_from_url():
     mock_response_text = """
     - name: CPU SU Rate
       history:
-        - value: 0.013
+        - value: "0.013"
           from: 2023-06
     """
     with requests_mock.Mocker() as m: