|
|
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) |
|
|
] |
|
|
|
