fix: Patch client to use in in DIRAC

Author: Robin VAN DE MERGHEL

diracx-client/src/diracx/client/_generated/aio/operations/_patch.py

@@ -11,10 +11,12 @@
 __all__ = [
     "AuthOperations",
     "JobsOperations",
+    "PilotsOperations"
 ]  # Add all objects you want publicly available to users at this package level
 
 from ....patches.auth.aio import AuthOperations
 from ....patches.jobs.aio import JobsOperations
+from ....patches.pilots.aio import PilotsOperations
 
 
 def patch_sdk():

diracx-client/src/diracx/client/_generated/operations/_patch.py

@@ -11,10 +11,12 @@
 __all__ = [
     "AuthOperations",
     "JobsOperations",
+    "PilotsOperations"
 ]  # Add all objects you want publicly available to users at this package level
 
 from ...patches.auth.sync import AuthOperations
 from ...patches.jobs.sync import JobsOperations
+from ...patches.pilots.sync import PilotsOperations
 
 
 def patch_sdk():

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

@@ -0,0 +1,34 @@
+"""Patches for the autorest-generated pilots client.
+
+This file can be used to customize the generated code for the pilots client.
+When adding new classes to this file, make sure to also add them to the
+__all__ list in the corresponding file in the patches directory.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "PilotsOperations",
+]
+
+from typing import Any, Unpack
+
+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
+
+# We're intentionally ignoring overrides here because we want to change the interface.
+# mypy: disable-error-code=override
+
+
+class PilotsOperations(_PilotsOperations):
+    @distributed_trace_async
+    async def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return await super().search(**make_search_body(**kwargs))
+
+    @distributed_trace_async
+    async def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return await super().summary(**make_summary_body(**kwargs))

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

@@ -0,0 +1,85 @@
+"""Utilities which are common to the sync and async pilots operator patches."""
+
+from __future__ import annotations
+
+__all__ = [
+    "make_search_body",
+    "SearchKwargs",
+    "make_summary_body",
+    "SummaryKwargs",
+]
+
+import json
+from io import BytesIO
+from typing import Any, IO, TypedDict, Unpack, cast, Literal
+
+from diracx.core.models import SearchSpec
+
+
+class ResponseExtra(TypedDict, total=False):
+    content_type: str
+    headers: dict[str, str]
+    params: dict[str, str]
+    cls: Any
+
+
+class SearchBody(TypedDict, total=False):
+    parameters: list[str] | None
+    search: list[SearchSpec] | None
+    sort: list[str] | None
+
+
+class SearchExtra(ResponseExtra, total=False):
+    page: int
+    per_page: int
+
+
+class SearchKwargs(SearchBody, SearchExtra): ...
+
+
+class UnderlyingSearchArgs(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_search_body(**kwargs: Unpack[SearchKwargs]) -> UnderlyingSearchArgs:
+    body: SearchBody = {}
+    for key in SearchBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["parameters", "search", "sort"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingSearchArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(SearchExtra, kwargs))
+    return result
+
+
+class SummaryBody(TypedDict, total=False):
+    grouping: list[str]
+    search: list[str]
+
+
+class SummaryKwargs(SummaryBody, ResponseExtra): ...
+
+
+class UnderlyingSummaryArgs(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_summary_body(**kwargs: Unpack[SummaryKwargs]) -> UnderlyingSummaryArgs:
+    body: SummaryBody = {}
+    for key in SummaryBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["grouping", "search"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingSummaryArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(ResponseExtra, kwargs))
+    return result

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

@@ -0,0 +1,34 @@
+"""Patches for the autorest-generated pilots client.
+
+This file can be used to customize the generated code for the pilots client.
+When adding new classes to this file, make sure to also add them to the
+__all__ list in the corresponding file in the patches directory.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "PilotsOperations",
+]
+
+from typing import Any, Unpack
+
+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
+
+# We're intentionally ignoring overrides here because we want to change the interface.
+# mypy: disable-error-code=override
+
+
+class PilotsOperations(_PilotsOperations):
+    @distributed_trace
+    def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return super().search(**make_search_body(**kwargs))
+
+    @distributed_trace
+    def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return super().summary(**make_summary_body(**kwargs))