pyproject, data_models: Add pandera & example for index data validation

mesmo/data_models/base_model.py

@@ -1,7 +1,10 @@
 """Custom pydantic base model."""
 
+from typing import Callable
+
 import numpy as np
 import pandas as pd
+import pandera as pa
 import pydantic as pyd
 from pydantic.json import timedelta_isoformat
 
@@ -19,3 +22,8 @@ class BaseModel(pyd.BaseModel):
             np.bool_: bool,
         },
     )
+
+
+def check_index(schema: pa.Index) -> Callable[[pd.Index], pd.Index]:
+    """Create pydantic validation function for pandas.Index based on given pandera.Index schema."""
+    return lambda v: schema.validate(pd.Series(index=v)).index

mesmo/data_models/model_index.py

@@ -1,10 +1,19 @@
 """Model index set data models."""
 
+from typing import Annotated
+
 import pandas as pd
+import pandera as pa
+import pydantic as pyd
 
 from mesmo.data_models import base_model
 
 
+class SampleModelIndex(base_model.BaseModel):
+    timesteps: Annotated[pd.Index, pyd.AfterValidator(base_model.check_index(pa.Index(pd.Timestamp)))]
+    names: Annotated[pd.Index, pyd.AfterValidator(base_model.check_index(pa.Index(str, unique=True)))]
+
+
 class ElectricGridModelIndex(base_model.BaseModel):
     timesteps: pd.Index
     phases: pd.Index

pyproject.toml

@@ -40,6 +40,7 @@ pyyaml = "^6.0.1"
 requests = "^2.31.0"                       # For HiGHS installation.
 scipy = "^1.11.3"
 tqdm = "^4.66.1"
+pandera = "^0.17.2"
 
 [tool.poetry.group.test.dependencies]
 coverage = { extras = ["toml"], version = "^7.3.2" }

tests/mesmo/data_models/model_index_test.py

@@ -0,0 +1,12 @@
+import pandas as pd
+import pytest
+
+from mesmo.data_models import model_index
+
+
+class TestSampleModelIndex:
+    def test_init(self):
+        timesteps = pd.Index([pd.Timestamp("now"), pd.Timestamp("now")])
+        names = pd.Index(["a", "b", "c"])
+
+        model_index.SampleModelIndex(timesteps=timesteps, names=names)