From 318666006de86c746ca96ee1bea22dee875c35d4 Mon Sep 17 00:00:00 2001
From: Ivan Ruiz Manuel <72193617+irm-codebase@users.noreply.github.com>
Date: Tue, 19 Nov 2024 17:17:58 +0100
Subject: [PATCH] Use Annotated for regex strings

---
 src/calliope/schemas/data_table_schema.py |  6 +----
 src/calliope/util/schema.py               |  5 ++--
 tests/test_data_table_schema.py           | 29 +++++++++++++++++++----
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/calliope/schemas/data_table_schema.py b/src/calliope/schemas/data_table_schema.py
index c693150f..0638d226 100644
--- a/src/calliope/schemas/data_table_schema.py
+++ b/src/calliope/schemas/data_table_schema.py
@@ -1,15 +1,11 @@
 """Implements the data table configuration class."""
 
-from typing import Self
-
 from pydantic import BaseModel, model_validator
+from typing_extensions import Self
 
 from calliope.util.schema import AttrStr, UniqueList
 from calliope.util.tools import listify
 
-# Get rid of pyright false negatives (see https://github.com/microsoft/pylance-release/issues/5457)
-# pyright: reportInvalidTypeForm=false
-
 
 class DataTable(BaseModel):
     """Data table validation model."""
diff --git a/src/calliope/util/schema.py b/src/calliope/util/schema.py
index f207c35a..f8cb7b6c 100644
--- a/src/calliope/util/schema.py
+++ b/src/calliope/util/schema.py
@@ -10,7 +10,7 @@
 from typing import Annotated, Literal, TypeVar
 
 import jsonschema
-from pydantic import AfterValidator, Field, constr
+from pydantic import AfterValidator, Field
 from pydantic_core import PydanticCustomError
 
 from calliope.attrdict import AttrDict
@@ -24,7 +24,8 @@
 
 # Regular string pattern for most calliope attributes
 FIELD_REGEX = r"^[^_^\d][\w]*$"
-AttrStr = constr(pattern=FIELD_REGEX)
+AttrStr = Annotated[str, Field(pattern=FIELD_REGEX)]
+
 # ==
 # Taken from https://github.com/pydantic/pydantic-core/pull/820#issuecomment-1670475909
 T = TypeVar("T", bound=Hashable)
diff --git a/tests/test_data_table_schema.py b/tests/test_data_table_schema.py
index 4e7b5c81..ba4d76b4 100644
--- a/tests/test_data_table_schema.py
+++ b/tests/test_data_table_schema.py
@@ -1,5 +1,7 @@
 """Test data table schema validation."""
 
+from pathlib import Path
+
 import pytest
 from pydantic import ValidationError
 
@@ -8,8 +10,10 @@
 
 from .common.util import check_error_or_warning
 
-FULL_TABLE_CONFIG = """
-data: time_varying_df
+
+@pytest.fixture
+def full_data_table_config():
+    return """data: time_varying_df
 rows: timesteps
 columns: [comment, nodes, techs]
 select:
@@ -25,6 +29,16 @@
 """
 
 
+@pytest.fixture
+def model_yaml_data_tables() -> AttrDict:
+    return AttrDict.from_yaml(
+        Path(__file__).parent
+        / "common"
+        / "national_scale_from_data_tables"
+        / "model.yaml"
+    )
+
+
 @pytest.mark.parametrize(
     "data_table",
     [{"rows": "timesteps"}, {"rows": "timesteps", "columns": ["techs", "nodes"]}],
@@ -71,7 +85,12 @@ def test_add_dims_overlap(rows, columns, add_dims):
     )
 
 
-@pytest.mark.parametrize("data_table", [FULL_TABLE_CONFIG])
-def test_full_table_config(data_table):
+def test_full_table_config(full_data_table_config):
     """Test a fully fledged data table configuration."""
-    DataTable(**AttrDict.from_yaml_string(data_table))
+    DataTable(**AttrDict.from_yaml_string(full_data_table_config))
+
+
+def test_data_table_model(model_yaml_data_tables):
+    """Data table validation must conform to expected usage."""
+    for data_table in model_yaml_data_tables["data_tables"].values():
+        DataTable(**data_table)