Annotation for defining Pydantic types based on AstroPy's Quantity

CONTRIBUTING.md

@@ -24,6 +24,46 @@ Facilities should attempt to avoid making excessive use of plain Python dictiona
 AEONlib also provides several [data types](https://github.com/AEONplus/AEONlib/blob/main/src/aeonlib/types.py) that improve Pydantic models:
 1. `aeonlib.types.Time` using this type in a Pydantic model allows consumers to pass in `astropy.time.Time` objects as well as `datetime` objects. The facility can then decide how the time is serialized to match whatever specific format is required.
 2. `aeonlib.types.Angle` similarly to the `Time` type, this allows consumers to pass in `astropy.coordinates.Angle` types as well as floats in decimal degrees, the facility can then decide how to serialize the type.
+3. `aeonlib.types.AstropyQuantityTypeAnnotation` can be used to define Pydantic models based on `astropy.units.Quantity`, which allow consumers to pass in values with units or as floats.
+
+The following example illustrates how these types can be used to define a model.
+
+```python
+from typing import Annotated, Union
+
+from astropy import coordinates
+from astropy import time
+from astropy import units as u
+from astropy.units import Quantity
+from pydantic import BaseModel
+
+from aeonlib.types import Angle, AstropyQuantityTypeAnnotation, Time
+
+Wavelength = Annotated[
+    Union[Quantity, float], AstropyQuantityTypeAnnotation(u.Angstrom)
+]
+
+class Observation(BaseModel):
+    start_time: Time
+    grating_angle: Angle
+    articulation_amgle: Angle
+    wavelength_range: tuple[Wavelength, Wavelength]
+
+observation = Observation(
+    start_time=time.Time(60775.0, scale="utc", format="mjd"),
+    grating_angle=22.5,
+    articulation_amgle=45 * u.deg,
+    wavelength_range=(5000, 600 * u.nm),  # 5000 Å to 6000 Å
+)
+
+assert type(observation.start_time) == time.Time
+assert type(observation.grating_angle) == coordinates.Angle
+assert type(observation.articulation_amgle) == coordinates.Angle
+assert type(observation.wavelength_range[0]) == Quantity
+assert type(observation.wavelength_range[1]) == Quantity
+```
+
+As you cannot use the validators of the `annotated_types` library with `Quantity` objects, the `aeonlib.validators` module provides some alternative validators, which you can use instead.
 
 These types eliminate the need for the facility user to need to remember which exact format a facility requires (time in hms? Or ISO UTC?) and simply pass in higher level objects instead.
 

src/aeonlib/types.py

@@ -1,12 +1,13 @@
 # pyright: reportUnknownVariableType=false
 # pyright: reportUnknownMemberType=false
+import dataclasses
 import logging
 from datetime import datetime
 from typing import Annotated, Any, cast
 
 import astropy.coordinates
 import astropy.time
-from astropy.units import Quantity
+from astropy.units import Quantity, UnitBase
 from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
 from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import core_schema
@@ -245,6 +246,90 @@ def __get_pydantic_json_schema__(
         }
 
 
+@dataclasses.dataclass
+class AstropyQuantityTypeAnnotation:
+    """
+    Annotation for defining custom Pydantic types based on a `astropy.units.Quantity`.
+
+    To define such a custom type, instantiate `AstropyQuantityTypeAnnotation` with
+    the default unit `d` and pass it as a type annotation. Pydantic fields with this
+    type can be instantiated wth a `float` or a `astropy.units.Quantity` with units
+    that are compatible with `d`. If a `float` is used, it is assumed to be given
+    with `d` as the unit. The field is stored as a `astropy.units.Quantity` with unit
+    `d`.
+
+    For example, you can define a ProperMotion type as follows:
+
+    ```
+    from typing import Annotated, Union
+    from astropy import units as u
+    from astropy.units import Quantity
+    from aeonlib.salt.models.types import AstropyQuantityTypeAnnotation
+
+    ProperMotion = Annotated[Union[Quantity, float], AstropyQuantityTypeAnnotation(u.arcsec / u.year)]
+    ```
+
+    This type can then be used in a Pydantic model:
+
+    ```
+    from pydantic import BaseModel
+
+    class MovingObject(BaseModel):
+        proper_motion: ProperMotion
+
+    # Create the same object in three different ways.
+    # Note: 1 year = 8766 hours
+    object1 = MovingObject(proper_motion=8766)  # 3 arcsec per year
+    object2 = MovingObject(proper_motion=8766 * u.arcsec / u.year)
+    object3 = MovingObject(proper_motion=1 * u.arcsec / u.hour)
+    """
+
+    # Based on
+    # https://docs.pydantic.dev/latest/concepts/types/#handling-third-party-types
+
+    default_unit: UnitBase
+
+    def __get_pydantic_core_schema__(
+        self,
+        _source_type: Any,
+        _handler: GetCoreSchemaHandler,
+    ) -> core_schema.CoreSchema:
+        def validate_from_float(value: float) -> Quantity:
+            return Quantity(value, unit=self.default_unit)
+
+        def validate_from_quantity(value: Quantity) -> Quantity:
+            return value.to(self.default_unit)
+
+        from_float_schema = core_schema.chain_schema(
+            [
+                core_schema.float_schema(),
+                core_schema.no_info_plain_validator_function(validate_from_float),
+            ]
+        )
+
+        from_quantity_schema = core_schema.chain_schema(
+            [
+                core_schema.is_instance_schema(Quantity),
+                core_schema.no_info_plain_validator_function(validate_from_quantity),
+            ]
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=from_float_schema,
+            python_schema=core_schema.union_schema(
+                [from_quantity_schema, from_float_schema]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda instance: instance.to(self.default_unit).value
+            ),
+        )
+
+    def __get_pydantic_json_schema(
+        self, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        return handler(core_schema.float_schema())
+
+
 Time = Annotated[astropy.time.Time | datetime, _AstropyTimeType]
 TimeMJD = Annotated[astropy.time.Time | datetime | float, _AstropyTimeMJDType]
 Angle = Annotated[

src/aeonlib/validators.py

@@ -0,0 +1,164 @@
+"""
+This module defines some Pydantic validators.
+
+The validators are
+"""
+
+from typing import Any
+
+import astropy.coordinates
+from astropy import units as u
+from pydantic import AfterValidator
+
+
+def _check_gt(a: Any, b: Any) -> Any:
+    if a <= b:
+        raise ValueError(f"{a} is not greater than to {b}.")
+    return a
+
+
+def _check_ge(a: Any, b: Any) -> Any:
+    if a < b:
+        raise ValueError(f"{a} is not greater than or equal to {b}.")
+    return a
+
+
+def _check_lt(a: Any, b: Any) -> None:
+    if a >= b:
+        raise ValueError(f"{a} is not less than to {b}.")
+    return a
+
+
+def _check_le(a: Any, b: Any) -> None:
+    if a > b:
+        raise ValueError(f"{a} is not less than or equal to {b}.")
+    return a
+
+
+def Gt(value: Any):
+    """
+    Return a Pydantic validator for checking a greater than relation.
+
+    The returned validator can be used in a type annotation::
+
+        import pydantic
+
+        class DummyModel(pydantic.BaseModel):
+            duration: Annotated[float, Gt(4)]
+
+    Pydantic will first perform its own internal validation and then check whether
+    the field value is greater than the argument passed to `Gt` (4 in the example
+    above).
+
+    It is up to the user to ensure that the field value and the argument of `Gt` can
+    be compared.
+
+    Parameters
+    ----------
+    value
+       Value against which to compare.
+
+    Returns
+    -------
+    A validator for checking a greater than relation.
+    """
+    return AfterValidator(lambda v: _check_gt(v, value))
+
+
+def Ge(value: Any):
+    """
+    Return a Pydantic validator for checking a greater than or equal to relation.
+
+    The returned validator can be used in a type annotation::
+
+        import pydantic
+
+        class DummyModel(pydantic.BaseModel):
+            duration: Annotated[float, Ge(4)]
+
+    Pydantic will first perform its own internal validation and then check whether
+    the field value is greater than or equal to the argument passed to `Ge` (4 in the
+    example above).
+
+    It is up to the user to ensure that the field value and the argument of `Ge` can
+    be compared.
+
+    Parameters
+    ----------
+    value
+       Value against which to compare.
+
+    Returns
+    -------
+    A validator for checking a greater than or equal to relation.
+    """
+    return AfterValidator(lambda v: _check_ge(v, value))
+
+
+def Lt(value: Any):
+    """
+    Return a Pydantic validator for checking a less than relation.
+
+    The returned validator can be used in a type annotation::
+
+        import pydantic
+
+        class DummyModel(pydantic.BaseModel):
+            height: Annotated[float, Lt(4)]
+
+    Pydantic will first perform its own internal validation and then check whether
+    the field value is less than or equal to the argument passed to `Lt` (4 in the
+    example above).
+
+    It is up to the user to ensure that the field value and the argument of `Lt` can
+    be compared.
+
+    Parameters
+    ----------
+    value
+       Value against which to compare.
+
+    Returns
+    -------
+    A validator for checking a less than relation.
+    """
+    return AfterValidator(lambda v: _check_lt(v, value))
+
+
+def Le(value: Any):
+    """
+    Return a Pydantic validator for checking a less than or equal to relation.
+
+    The returned validator can be used in a type annotation::
+
+        import pydantic
+
+        class DummyModel(pydantic.BaseModel):
+            height: Annotated[float, Le(4)]
+
+    Pydantic will first perform its own internal validation and then check whether
+    the field value is less than or equal to the argument passed to `Le` (4 in the
+    example above).
+
+    It is up to the user to ensure that the field value and the argument of `Le` can
+    be compared.
+
+    Parameters
+    ----------
+    value
+       Value against which to compare.
+
+    Returns
+    -------
+    A validator for checking a less than or equal to relation.
+    """
+    return AfterValidator(lambda v: _check_le(v, value))
+
+
+def check_in_visibility_range(
+    dec: astropy.coordinates.Angle,
+) -> astropy.coordinates.Angle:
+    if dec < -76 * u.deg or dec > 11 * u.deg:
+        raise ValueError("Not in SALT's visibility range (between -76 and 11 degrees).")
+
+    return dec