fix: Add patch for the pilot client

Author: Robin-Van-de-Merghel

diracx-client/src/diracx/client/patches/pilots/aio.py

@@ -16,7 +16,16 @@
 from azure.core.tracing.decorator_async import distributed_trace_async
 
 from ..._generated.aio.operations._operations import PilotsOperations as _PilotsOperations
-from .common import make_search_body, make_summary_body, SearchKwargs, SummaryKwargs
+from .common import (
+    make_search_body,
+    make_summary_body,
+    make_add_pilot_stamps_body,
+    make_update_pilot_fields_body,
+    SearchKwargs,
+    SummaryKwargs,
+    AddPilotStampsKwargs,
+    UpdatePilotFieldsKwargs
+)
 
 # We're intentionally ignoring overrides here because we want to change the interface.
 # mypy: disable-error-code=override
@@ -32,3 +41,13 @@ async def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
     async def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
         """TODO"""
         return await super().summary(**make_summary_body(**kwargs))
+
+    @distributed_trace_async
+    async def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> None:
+        """TODO"""
+        return await super().add_pilot_stamps(**make_add_pilot_stamps_body(**kwargs))
+
+    @distributed_trace_async
+    async def update_pilot_fields(self, **kwargs: Unpack[UpdatePilotFieldsKwargs]) -> None:
+        """TODO"""
+        return await super().update_pilot_fields(**make_update_pilot_fields_body(**kwargs))

diracx-client/src/diracx/client/patches/pilots/common.py

@@ -7,13 +7,17 @@
     "SearchKwargs",
     "make_summary_body",
     "SummaryKwargs",
+    "AddPilotStampsKwargs",
+    "make_add_pilot_stamps_body",
+    "UpdatePilotFieldsKwargs",
+    "make_update_pilot_fields_body"
 ]
 
 import json
 from io import BytesIO
 from typing import Any, IO, TypedDict, Unpack, cast, Literal
 
-from diracx.core.models import SearchSpec
+from diracx.core.models import SearchSpec, PilotStatus, PilotFieldsMapping
 
 
 class ResponseExtra(TypedDict, total=False):
@@ -23,6 +27,7 @@ class ResponseExtra(TypedDict, total=False):
     cls: Any
 
 
+# ------------------ Search ------------------
 class SearchBody(TypedDict, total=False):
     parameters: list[str] | None
     search: list[SearchSpec] | None
@@ -56,6 +61,7 @@ def make_search_body(**kwargs: Unpack[SearchKwargs]) -> UnderlyingSearchArgs:
     result.update(cast(SearchExtra, kwargs))
     return result
 
+# ------------------ Summary ------------------
 
 class SummaryBody(TypedDict, total=False):
     grouping: list[str]
@@ -83,3 +89,57 @@ def make_summary_body(**kwargs: Unpack[SummaryKwargs]) -> UnderlyingSummaryArgs:
     result: UnderlyingSummaryArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
     result.update(cast(ResponseExtra, kwargs))
     return result
+
+# ------------------ AddPilotStamps ------------------
+
+class AddPilotStampsBody(TypedDict, total=False):
+    pilot_stamps: list[str]
+    grid_type: str
+    grid_site: str
+    pilot_references: dict[str, str]
+    pilot_status: PilotStatus
+
+class AddPilotStampsKwargs(AddPilotStampsBody, ResponseExtra): ...
+
+class UnderlyingAddPilotStampsArgs(ResponseExtra, total=False):
+    # FIXME: The autorest-generated has a bug that it expected IO[bytes] despite
+    # the code being generated to support IO[bytes] | bytes.
+    body: IO[bytes]
+
+def make_add_pilot_stamps_body(**kwargs: Unpack[AddPilotStampsKwargs]) -> UnderlyingAddPilotStampsArgs:
+    body: AddPilotStampsBody = {}
+    for key in AddPilotStampsBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["pilot_stamps", "grid_type", "grid_site", "pilot_references", "pilot_status"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingAddPilotStampsArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(ResponseExtra, kwargs))
+    return result
+
+# ------------------ UpdatePilotFields ------------------
+
+class UpdatePilotFieldsBody(TypedDict, total=False):
+    pilot_stamps_to_fields_mapping: list[PilotFieldsMapping]
+
+class UpdatePilotFieldsKwargs(UpdatePilotFieldsBody, ResponseExtra): ...
+
+class UnderlyingUpdatePilotFields(ResponseExtra, total=False):
+    # FIXME: The autorest-generated has a bug that it expected IO[bytes] despite
+    # the code being generated to support IO[bytes] | bytes.
+    body: IO[bytes]
+
+def make_update_pilot_fields_body(**kwargs: Unpack[UpdatePilotFieldsKwargs]) -> UnderlyingUpdatePilotFields:
+    body: UpdatePilotFieldsBody = {}
+    for key in UpdatePilotFieldsBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["pilot_stamps_to_fields_mapping"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingUpdatePilotFields = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(ResponseExtra, kwargs))
+    return result

diracx-client/src/diracx/client/patches/pilots/sync.py

@@ -16,7 +16,16 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from ..._generated.operations._operations import PilotsOperations as _PilotsOperations
-from .common import make_search_body, make_summary_body, SearchKwargs, SummaryKwargs
+from .common import (
+    make_search_body,
+    make_summary_body,
+    make_add_pilot_stamps_body,
+    make_update_pilot_fields_body,
+    SearchKwargs,
+    SummaryKwargs,
+    AddPilotStampsKwargs,
+    UpdatePilotFieldsKwargs
+)
 
 # We're intentionally ignoring overrides here because we want to change the interface.
 # mypy: disable-error-code=override
@@ -32,3 +41,13 @@ def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
     def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
         """TODO"""
         return super().summary(**make_summary_body(**kwargs))
+
+    @distributed_trace
+    def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> None:
+        """TODO"""
+        return super().add_pilot_stamps(**make_add_pilot_stamps_body(**kwargs))
+
+    @distributed_trace
+    def update_pilot_fields(self, **kwargs: Unpack[UpdatePilotFieldsKwargs]) -> None:
+        """TODO"""
+        return super().update_pilot_fields(**make_update_pilot_fields_body(**kwargs))