Merge pull request #43 from jenskeiner/feature/defaults-examples-and-…

…schemas Add JSON schema and examples to custom types.
jenskeiner · May 23, 2024 · 0a3282d · 0a3282d
2 parents ad0d5b0 + d459a06
commit 0a3282d
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 20 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,4 +1,9 @@
 repos:
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.15.2
+    hooks:
+      - id: pyupgrade
+        args: ["--py39-plus"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
     rev: v0.3.2
@@ -7,8 +12,3 @@ repos:
       - id: ruff
       # Run the formatter.
       - id: ruff-format
-  - repo: https://github.com/asottile/pyupgrade
-    rev: v3.15.2
-    hooks:
-      - id: pyupgrade
-        args: ["--py39-plus"]
diff --git a/exchange_calendars_extensions/api/changes.py b/exchange_calendars_extensions/api/changes.py
@@ -5,8 +5,16 @@
 from typing import Union, Annotated, Callable
 
 import pandas as pd
-from pydantic import BaseModel, Field, RootModel, model_validator, validate_call
-from pydantic.functional_validators import BeforeValidator, AfterValidator
+from pydantic import (
+    BaseModel,
+    Field,
+    RootModel,
+    model_validator,
+    validate_call,
+    WithJsonSchema,
+    BeforeValidator,
+    AfterValidator,
+)
 from typing_extensions import (
     Literal,
     Any,
@@ -100,14 +108,24 @@ def _to_date(value: pd.Timestamp) -> pd.Timestamp:
 
 
 # A type alias for pd.Timestamp that allows initialisation from suitably formatted string values.
-TimestampLike = Annotated[pd.Timestamp, BeforeValidator(_to_timestamp)]
+TimestampLike = Annotated[
+    pd.Timestamp,
+    BeforeValidator(_to_timestamp),
+    WithJsonSchema({"type": "string", "format": "date-time"}),
+    Field(examples=["2020-01-01T00:00:00Z"]),
+]
 
 # A type alias for TimestampLike that normalizes the timestamp to a date-like value.
 #
 # Date-like means that the timestamp is timezone-naive and normalized to the date boundary, i.e. midnight of the day it
 # represents. If the input converts to a valid pd.Timestamp, any timezone information, if present, is discarded. If the
 # result is not aligned with a date boundary, it is normalized to midnight of the same day.
-DateLike = Annotated[TimestampLike, AfterValidator(_to_date)]
+DateLike = Annotated[
+    TimestampLike,
+    AfterValidator(_to_date),
+    WithJsonSchema({"type": "string", "format": "date"}),
+    Field(examples=["2020-01-01"]),
+]
 
 
 class AbstractDayProps(
@@ -117,17 +135,21 @@ class AbstractDayProps(
     Abstract base class for special day properties.
     """
 
-    name: str  # The name of the day.
+    name: str = Field(
+        examples=["Holiday", "Ad-hoc Holiday", "Special Close Day", "Special Open Day"]
+    )  # The name of the day.
 
 
 class DayProps(AbstractDayProps):
     """
     Vanilla special day specification.
     """
 
-    type: Literal[
-        DayType.HOLIDAY, DayType.MONTHLY_EXPIRY, DayType.QUARTERLY_EXPIRY
-    ]  # The type of the special day.
+    type: Literal[DayType.HOLIDAY, DayType.MONTHLY_EXPIRY, DayType.QUARTERLY_EXPIRY] = (
+        Field(
+            examples=[DayType.HOLIDAY, DayType.MONTHLY_EXPIRY, DayType.QUARTERLY_EXPIRY]
+        )
+    )  # The type of the special day.
 
     def __str__(self):
         return f'{{type={self.type.name}, name="{self.name}"}}'
@@ -168,7 +190,12 @@ def _to_time(value: Union[dt.time, str]):
 
 
 # A type alias for dt.time that allows initialisation from suitably formatted string values.
-TimeLike = Annotated[dt.time, BeforeValidator(_to_time)]
+TimeLike = Annotated[
+    dt.time,
+    BeforeValidator(_to_time),
+    WithJsonSchema({"type": "string", "format": "time"}),
+    Field(examples=["09:00", "16:30"]),
+]
 
 
 class DayPropsWithTime(AbstractDayProps):
@@ -188,7 +215,10 @@ def __str__(self):
 # Type alias for valid day properties.
 DayPropsLike = Annotated[Union[DayProps, DayPropsWithTime], Field(discriminator="type")]
 
-Tags = Union[list[str], Union[tuple[str], Union[set[str], None]]]
+Tags = Annotated[
+    Union[list[str], Union[tuple[str], Union[set[str], None]]],
+    Field(examples=[["tag1", "tag2"]]),
+]
 
 
 class DayMeta(
@@ -202,7 +232,7 @@ class DayMeta(
     tags: Tags = []
 
     # Free-form comment.
-    comment: Union[str, None] = None
+    comment: Union[str, None] = Field(default=None, examples=["This is a comment."])
 
     @model_validator(mode="after")
     def _canonicalize(self) -> "DayMeta":
@@ -296,9 +326,17 @@ class ChangeSet(
     already exist in any day type category, and updated/moved to the new day type if it does.
     """
 
-    add: dict[DateLike, DayPropsLike] = Field(default_factory=dict)
-    remove: list[DateLike] = Field(default_factory=list)
-    meta: dict[DateLike, DayMeta] = Field(default_factory=dict)
+    add: dict[DateLike, DayPropsLike] = Field(
+        default_factory=dict,
+        examples=[{"2020-01-01": {"type": "holiday", "name": "New Year's Day"}}],
+    )
+    remove: list[DateLike] = Field(default_factory=list, examples=["2020-01-01"])
+    meta: dict[DateLike, DayMeta] = Field(
+        default_factory=dict,
+        examples=[
+            {"2020-01-01": {"tags": ["tag1", "tag2"], "comment": "This is a comment."}}
+        ],
+    )
 
     @model_validator(mode="after")
     def _canonicalize(self) -> "ChangeSet":
@@ -552,7 +590,7 @@ def all_days(self, include_meta: bool = False) -> tuple[pd.Timestamp, ...]:
 
 # A type alias for a dictionary of changesets, mapping exchange key to a corresponding change set.
 class ChangeSetDict(RootModel):
-    root: dict[str, ChangeSet]
+    root: dict[str, ChangeSet] = Field(default_factory=dict)
 
     # Delegate all dictionary-typical methods to the root dictionary.
     def __getitem__(self, key):