File size: 2,311 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from __future__ import annotations
from itertools import chain
from typing import TYPE_CHECKING, Any, Iterable, Mapping
from pydantic import ValidationError
from typing_extensions import override
from wandb_graphql.language.ast import Document
from wandb.apis.paginator import Paginator, _Client
if TYPE_CHECKING:
from wandb.automations import Automation
from wandb.automations._generated import ProjectConnectionFields
class Automations(Paginator["Automation"]):
last_response: ProjectConnectionFields | None
_query: Document
def __init__(
self,
client: _Client,
variables: Mapping[str, Any],
per_page: int = 50,
_query: Document | None = None,
):
super().__init__(client, variables, per_page=per_page)
if _query is None:
raise RuntimeError(f"Query required for {type(self).__qualname__}")
self._query = _query
@property
def more(self) -> bool:
"""Whether there are more items to fetch."""
if self.last_response is None:
return True
return self.last_response.page_info.has_next_page
@property
def cursor(self) -> str | None:
"""The start cursor to use for the next page."""
if self.last_response is None:
return None
return self.last_response.page_info.end_cursor
@override
def _update_response(self) -> None:
"""Fetch the raw response data for the current page."""
from wandb.automations._generated import ProjectConnectionFields
data: dict[str, Any] = self.client.execute(
self._query, variable_values=self.variables
)
try:
page_data = data["searchScope"]["projects"]
self.last_response = ProjectConnectionFields.model_validate(page_data)
except (LookupError, AttributeError, ValidationError) as e:
raise ValueError("Unexpected response data") from e
def convert_objects(self) -> Iterable[Automation]:
"""Parse the page data into a list of objects."""
from wandb.automations import Automation
page = self.last_response
return [
Automation.model_validate(obj)
for obj in chain.from_iterable(edge.node.triggers for edge in page.edges)
]
|