DM-44763: Take advantage of class serialization

src/vocutouts/models/parameters.py

@@ -7,7 +7,7 @@
 
 from ..exceptions import InvalidCutoutParameterError
 from ..uws.models import UWSJobParameter
-from .stencils import Stencil, parse_stencil
+from .stencils import CircleStencil, PolygonStencil, RangeStencil, Stencil
 
 
 @dataclass
@@ -46,8 +46,9 @@ def from_job_parameters(cls, params: list[UWSJobParameter]) -> Self:
                 if param.parameter_id == "id":
                     ids.append(param.value)
                 else:
-                    f = parse_stencil(param.parameter_id.upper(), param.value)
-                    stencils.append(f)
+                    stencil_type = param.parameter_id.upper()
+                    stencil = cls._parse_stencil(stencil_type, param.value)
+                    stencils.append(stencil)
         except Exception as e:
             msg = f"Invalid cutout parameter: {type(e).__name__}: {e!s}"
             raise InvalidCutoutParameterError(msg, params) from e
@@ -58,3 +59,18 @@ def from_job_parameters(cls, params: list[UWSJobParameter]) -> Self:
                 "No cutout stencil given", params
             )
         return cls(ids=ids, stencils=stencils)
+
+    @staticmethod
+    def _parse_stencil(stencil_type: str, params: str) -> Stencil:
+        """Convert a string stencil parameter to its representation."""
+        if stencil_type == "POS":
+            stencil_type, params = params.split(None, 1)
+        match stencil_type:
+            case "CIRCLE":
+                return CircleStencil.from_string(params)
+            case "POLYGON":
+                return PolygonStencil.from_string(params)
+            case "RANGE":
+                return RangeStencil.from_string(params)
+            case _:
+                raise ValueError(f"Unknown stencil type {stencil_type}")

src/vocutouts/models/stencils.py

@@ -4,12 +4,21 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Self
+from typing import Any, Self, TypeAlias
 
 from astropy import units as u
 from astropy.coordinates import Angle, SkyCoord
 
-Range = tuple[float, float]
+Range: TypeAlias = tuple[float, float]
+"""Type representing a range of a coordinate."""
+
+__all__ = [
+    "CircleStencil",
+    "PolygonStencil",
+    "Range",
+    "RangeStencil",
+    "Stencil",
+]
 
 
 class Stencil(ABC):
@@ -22,15 +31,18 @@ def from_string(cls, params: str) -> Self:
 
     @abstractmethod
     def to_dict(self) -> dict[str, Any]:
-        """Convert the stencil to a JSON-serializable form for queuing."""
+        """Convert the stencil to a dictionary for logging."""
 
 
 @dataclass
 class CircleStencil(Stencil):
     """Represents a ``CIRCLE`` or ``POS=CIRCLE`` stencil."""
 
     center: SkyCoord
+    """Center of the circle."""
+
     radius: Angle
+    """Radius of the circle."""
 
     @classmethod
     def from_string(cls, params: str) -> Self:
@@ -60,6 +72,7 @@ class PolygonStencil(Stencil):
     """
 
     vertices: SkyCoord
+    """Vertices of the polygon."""
 
     @classmethod
     def from_string(cls, params: str) -> Self:
@@ -90,7 +103,10 @@ class RangeStencil(Stencil):
     """Represents a ``POS=RANGE`` stencil."""
 
     ra: Range
+    """Range of ra values."""
+
     dec: Range
+    """Range of dec values."""
 
     @classmethod
     def from_string(cls, params: str) -> Self:
@@ -106,17 +122,3 @@ def to_dict(self) -> dict[str, Any]:
             "ra": self.ra,
             "dec": self.dec,
         }
-
-
-def parse_stencil(stencil_type: str, params: str) -> Stencil:
-    """Convert a string stencil parameter to its internal representation."""
-    if stencil_type == "POS":
-        stencil_type, params = params.split(None, 1)
-    if stencil_type == "CIRCLE":
-        return CircleStencil.from_string(params)
-    elif stencil_type == "POLYGON":
-        return PolygonStencil.from_string(params)
-    elif stencil_type == "RANGE":
-        return RangeStencil.from_string(params)
-    else:
-        raise ValueError(f"Unknown stencil type {stencil_type}")

src/vocutouts/policy.py

@@ -33,16 +33,16 @@ def __init__(self, arq: ArqQueue, logger: BoundLogger) -> None:
         super().__init__(arq)
         self._logger = logger
 
-    async def dispatch(self, job: UWSJob, access_token: str) -> JobMetadata:
+    async def dispatch(self, job: UWSJob, token: str) -> JobMetadata:
         """Dispatch a cutout request to the backend.
 
         Parameters
         ----------
         job
             The submitted job description.
-        access_token
-            Gafaelfawr access token used to authenticate to Butler server
-            in the backend.
+        token
+            Gafaelfawr token used to authenticate to the Butler server in the
+            backend.
 
         Returns
         -------
@@ -54,13 +54,9 @@ async def dispatch(self, job: UWSJob, access_token: str) -> JobMetadata:
         Currently, only one dataset ID and only one stencil are supported.
         This limitation is expected to be relaxed in a later version.
         """
-        cutout_params = CutoutParameters.from_job_parameters(job.parameters)
+        params = CutoutParameters.from_job_parameters(job.parameters)
         return await self.arq.enqueue(
-            "cutout",
-            job.job_id,
-            cutout_params.ids,
-            [s.to_dict() for s in cutout_params.stencils],
-            access_token,
+            "cutout", job.job_id, params.ids, params.stencils, token=token
         )
 
     def validate_params(self, params: list[UWSJobParameter]) -> None:

src/vocutouts/uws/policy.py

@@ -40,7 +40,7 @@ def __init__(self, arq: ArqQueue) -> None:
         self.arq = arq
 
     @abstractmethod
-    async def dispatch(self, job: UWSJob, access_token: str) -> JobMetadata:
+    async def dispatch(self, job: UWSJob, token: str) -> JobMetadata:
         """Dispatch a job to a backend worker.
 
         This method is responsible for converting UWS job parameters to the
@@ -51,9 +51,9 @@ async def dispatch(self, job: UWSJob, access_token: str) -> JobMetadata:
         ----------
         job
             Job to start.
-        access_token
-            Gafaelfawr access token used to authenticate to services used
-            by the backend on the user's behalf.
+        token
+            Gafaelfawr token used to authenticate to services used by the
+            backend on the user's behalf.
 
         Returns
         -------

src/vocutouts/uws/service.py

@@ -226,9 +226,7 @@ async def list_jobs(
             user, phases=phases, after=after, count=count
         )
 
-    async def start(
-        self, user: str, job_id: str, access_token: str
-    ) -> JobMetadata:
+    async def start(self, user: str, job_id: str, token: str) -> JobMetadata:
         """Start execution of a job.
 
         Parameters
@@ -237,9 +235,9 @@ async def start(
             User on behalf of whom this operation is performed.
         job_id
             Identifier of the job to start.
-        access_token
-            Gafaelfawr access token used to authenticate to services used
-            by the backend on the user's behalf.
+        token
+            Gafaelfawr token used to authenticate to services used by the
+            backend on the user's behalf.
 
         Returns
         -------
@@ -257,7 +255,7 @@ async def start(
             raise PermissionDeniedError(f"Access to job {job_id} denied")
         if job.phase not in (ExecutionPhase.PENDING, ExecutionPhase.HELD):
             raise InvalidPhaseError("Cannot start job in phase {job.phase}")
-        metadata = await self._policy.dispatch(job, access_token)
+        metadata = await self._policy.dispatch(job, token)
         await self._storage.mark_queued(job_id, metadata)
         return metadata
 

src/vocutouts/workers/cutout.py

@@ -8,13 +8,10 @@
 
 import os
 from datetime import timedelta
-from typing import Any
 from urllib.parse import urlparse
 from uuid import UUID
 
-import astropy.units as u
 import structlog
-from astropy.coordinates import Angle, SkyCoord
 from lsst.afw.geom import SinglePolygonException
 from lsst.daf.butler import LabeledButlerFactory
 from lsst.image_cutout_backend import ImageCutoutBackend, projection_finders
@@ -23,6 +20,7 @@
 from safir.logging import configure_logging
 from structlog.stdlib import BoundLogger
 
+from ..models.stencils import CircleStencil, PolygonStencil, Stencil
 from ..uws.exceptions import TaskFatalError, TaskUserError
 from ..uws.models import ErrorCode, UWSJobResult
 from ..uws.workers import UWSWorkerConfig, build_worker
@@ -33,7 +31,7 @@
 __all__ = ["WorkerSettings"]
 
 
-def _get_backend(butler_label: str, access_token: str) -> ImageCutoutBackend:
+def _get_backend(butler_label: str, token: str) -> ImageCutoutBackend:
     """Given the Butler label, retrieve or build a backend.
 
     The dataset ID will be a URI of the form ``butler://<label>/<uuid>``.
@@ -44,16 +42,16 @@ def _get_backend(butler_label: str, access_token: str) -> ImageCutoutBackend:
     ----------
     butler_label
         Label portion of the Butler URI.
-    access_token
-        Gafaelfawr access token, used to access the Butler service.
+    token
+        Gafaelfawr token, used to access the Butler service.
 
     Returns
     -------
     lsst.image_cutout_backend.ImageCutoutBackend
         Backend to use.
     """
     butler = _BUTLER_FACTORY.create_butler(
-        label=butler_label, access_token=access_token
+        label=butler_label, access_token=token
     )
 
     # At present, projection finders and image cutout backend have no internal
@@ -87,9 +85,9 @@ def _parse_uri(uri: str) -> tuple[str, UUID]:
 def cutout(
     job_id: str,
     dataset_ids: list[str],
-    stencils: list[dict[str, Any]],
-    access_token: str,
+    stencils: list[Stencil],
     *,
+    token: str,
     logger: BoundLogger,
 ) -> list[UWSJobResult]:
     """Perform a cutout.
@@ -103,17 +101,15 @@ def cutout(
     Parameters
     ----------
     job_id
-        UWS job ID, used as the key for storing results.
+        UWS job ID, provided by the UWS layer but not used here.
     dataset_ids
         Data objects on which to perform cutouts. These are opaque identifiers
         passed as-is to the backend. The user will normally discover them via
         some service such as ObsTAP.
     stencils
-        Serialized stencils for the cutouts to perform. These are
-        `~vocutouts.models.stencils.Stencil` objects corresponding to the
-        user's request.
-    access_token
-        Gafaelfawr access token used to authenticate to Butler server.
+        Serialized stencils for the cutouts to perform.
+    token
+        Gafaelfawr token used to authenticate to Butler server.
     logger
         Logger to use for logging.
 
@@ -131,7 +127,9 @@ def cutout(
         Raised if the cutout failed for reasons that may go away if the cutout
         is retried.
     """
-    logger = logger.bind(dataset_ids=dataset_ids, stencils=stencils)
+    logger = logger.bind(
+        dataset_ids=dataset_ids, stencils=[s.to_dict() for s in stencils]
+    )
 
     # Currently, only a single dataset ID and a single stencil are supported.
     # These constraints should have been applied by the policy layer, so if we
@@ -145,29 +143,21 @@ def cutout(
 
     # Parse the dataset ID and retrieve an appropriate backend.
     butler_label, uuid = _parse_uri(dataset_ids[0])
-    backend = _get_backend(butler_label, access_token)
+    backend = _get_backend(butler_label, token)
 
     # Convert the stencils to SkyStencils.
     sky_stencils = []
-    for stencil_dict in stencils:
-        if stencil_dict["type"] == "circle":
-            center = SkyCoord(
-                stencil_dict["center"]["ra"] * u.degree,
-                stencil_dict["center"]["dec"] * u.degree,
-                frame="icrs",
-            )
-            radius = Angle(stencil_dict["radius"] * u.degree)
-            stencil = SkyCircle.from_astropy(center, radius, clip=True)
-        elif stencil_dict["type"] == "polygon":
-            ras = [v[0] for v in stencil_dict["vertices"]]
-            decs = [v[1] for v in stencil_dict["vertices"]]
-            vertices = SkyCoord(ras * u.degree, decs * u.degree, frame="icrs")
-            stencil = SkyPolygon.from_astropy(vertices, clip=True)
-        else:
-            msg = f'Unknown stencil type {stencil_dict["type"]}'
-            logger.warning(msg)
-            raise TaskFatalError(ErrorCode.USAGE_ERROR, msg)
-        sky_stencils.append(stencil)
+    for stencil in stencils:
+        match stencil:
+            case CircleStencil(center, radius):
+                sky_stencil = SkyCircle.from_astropy(center, radius, clip=True)
+            case PolygonStencil(vertices):
+                sky_stencil = SkyPolygon.from_astropy(vertices, clip=True)
+            case _:
+                msg = f"Unknown stencil type {type(stencil).__name__}"
+                logger.warning(msg)
+                raise TaskFatalError(ErrorCode.USAGE_ERROR, msg)
+        sky_stencils.append(sky_stencil)
 
     # Perform the cutout. We have no idea if unknown exceptions here are
     # transient or fatal, so conservatively assume they are fatal. Provide a

tests/models/parameters_test.py

@@ -0,0 +1,28 @@
+"""Tests for cutout parameter models."""
+
+from __future__ import annotations
+
+import pickle
+
+from vocutouts.models.parameters import CutoutParameters
+from vocutouts.models.stencils import (
+    CircleStencil,
+    PolygonStencil,
+    RangeStencil,
+)
+
+
+def test_pickle() -> None:
+    params = CutoutParameters(
+        ids=["foo", "bar"],
+        stencils=[
+            CircleStencil.from_string("1 1.42 1"),
+            PolygonStencil.from_string("1 0 1 1 0 1 0 0"),
+            RangeStencil.from_string("-Inf 1 0 1"),
+        ],
+    )
+    params_pickle = pickle.loads(pickle.dumps(params))
+    assert params.ids == params_pickle.ids
+    expected = [s.to_dict() for s in params.stencils]
+    seen = [s.to_dict() for s in params_pickle.stencils]
+    assert expected == seen