feat: an expressions entry built in code has its keys checked

src/math_spec/__init__.py

@@ -32,6 +32,8 @@
     VARIABLE_DOMAINS,
     SosBlock,
     Spec,
+    WrittenCase,
+    WrittenExpression,
 )
 from math_spec.operators import (
     BUILTIN_NAMES,
@@ -70,6 +72,8 @@
     'SosBlock',
     'Spec',
     'SymbolTable',
+    'WrittenCase',
+    'WrittenExpression',
     'advice',
     'call_shape_error',
     'did_you_mean',

src/math_spec/model.py

@@ -13,7 +13,7 @@
 import re
 from collections import Counter
 from functools import cached_property
-from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, Self, cast, get_args, override
+from typing import TYPE_CHECKING, Annotated, ClassVar, Literal, Self, TypedDict, cast, get_args, override
 
 from pydantic import (
     BaseModel,
@@ -476,6 +476,40 @@ def _as_written(self) -> str | dict[str, object]:
         return {'expression': self.expression, 'description': self.description}
 
 
+class WrittenCase(TypedDict):
+    """One entry under ``cases:``, in the shape a file writes it.
+
+    Both keys are required, as :class:`ExpressionCase` requires them.
+    """
+
+    when: str
+    expression: str | float
+
+
+class WrittenExpression(TypedDict, total=False):
+    """One ``expressions:`` entry as a mapping, in the shape a file writes it.
+
+    For a caller that builds the mapping in code rather than reading it from
+    YAML: the keys are checked where the caller writes them, instead of at
+    :func:`~math_spec.to_spec`.
+
+    Every key is optional, because which *combination* is a model is
+    :class:`ExpressionBlock`'s own rule — one ``expression:``, or ``cases:``
+    with the ``dims:`` and ``otherwise:`` they need. This says which keys
+    exist and what each takes; loading still decides the rest, and still
+    refuses a key that is not here.
+
+    Example:
+        >>> written: WrittenExpression = {'expression': 'sum(p, over=generator)'}
+    """
+
+    expression: str | float
+    dims: list[str]
+    cases: dict[str, WrittenCase]
+    otherwise: str | float
+    description: str
+
+
 class PiecewiseLink(_StrictBlock):
     """One link of a piecewise block: an expression pinned to a values curve.
 

tests/test_public_surface.py

@@ -36,6 +36,8 @@
         # the declaration vocabularies a consumer pins its own tables against
         'DIMENSION_DTYPES', 'PARAMETER_DTYPES', 'VARIABLE_DOMAINS', 'VARIABLE_ABSENCE', 'ADVICE_KINDS',
         'CURVATURES', 'SosBlock',
+        # what a caller writing an expressions: entry in code annotates it with
+        'WrittenExpression', 'WrittenCase',
         # typesetting
         'FORMATS', 'SymbolTable', 'typeset', 'typeset_declaration', 'to_latex', 'to_typst', 'to_markdown',
     }

tests/test_schema.py

@@ -2,12 +2,16 @@
 #
 # SPDX-License-Identifier: MIT
 
-"""The published JSON Schema is the pydantic models, verbatim.
+"""What the package publishes about the YAML surface is the pydantic models, verbatim.
 
 `schema/math-spec.schema.json` is a generated artefact that ships in the
 repository so an editor can offer completion without importing the package.
 Nothing regenerates it on the way to a release, so the only thing keeping it
 equal to the models is this file.
+
+`WrittenExpression` and `WrittenCase` are the same surface a second time, for
+a caller that builds an `expressions:` entry in code. They are written by hand
+rather than generated, so the same job falls here.
 """
 
 import json
@@ -79,3 +83,25 @@ def test_the_piecewise_method_vocabulary_has_one_home():
     assert set(get_args(model.PiecewiseMethod)) == set(model.PIECEWISE_METHODS), (
         'the typed methods and the emitting ones disagree, so a method is accepted that emits nothing or the reverse'
     )
+
+
+@pytest.mark.parametrize(
+    ('written', 'definition'),
+    [
+        pytest.param(model.WrittenExpression, 'ExpressionBlock', id='expression'),
+        pytest.param(model.WrittenCase, 'ExpressionCase', id='case'),
+    ],
+)
+def test_the_written_form_takes_the_keys_its_block_takes(written, definition):
+    """A field added to the block and not to the TypedDict would leave a caller
+    annotating the key it just gained as an error, and one removed would leave
+    them annotating a key the loader now refuses. Read off the schema rather
+    than the model, because that is the rendering the shorthand forms reach."""
+    published = json.loads(schema.PATH.read_text())['$defs'][definition]
+    mapping = next(form for form in published.get('anyOf', [published]) if 'properties' in form)
+    assert set(written.__annotations__) == set(mapping['properties']), (
+        f'{written.__name__} and {definition} no longer take the same keys'
+    )
+    assert written.__required_keys__ == frozenset(mapping.get('required', ())), (
+        f'{written.__name__} and {definition} disagree about which keys are required'
+    )