[IO-1279] Functions to move items around stages (#640)

* WIP: moving items to stage

* WIP: moving items to stage

* WIP: moving items to stage

* pushing items between stages

* added stage syntax

* WIP dogfooding functions

* fixing broken tests

* remove curses

* tests for query

* tests for core

* tests for stageMeta

* dogfooding feedback

* dogfooding feedback

* fixes for tests

darwin/future/core/client.py

@@ -155,7 +155,7 @@ def _setup_session(self, retries: Retry) -> None:
 
     @property
     def headers(self) -> Dict[str, str]:
-        http_headers: Dict[str, str] = {"Content-Type": "application/json"}
+        http_headers: Dict[str, str] = {"Content-Type": "application/json", "Accept": "application/json"}
         if self.config.api_key:
             http_headers["Authorization"] = f"ApiKey {self.config.api_key}"
         return http_headers
@@ -172,7 +172,7 @@ def _generic_call(self, method: Callable, endpoint: str, payload: Optional[dict]
         endpoint = self._sanitize_endpoint(endpoint)
         url = self.config.api_endpoint + endpoint
         if payload is not None:
-            response = method(url, payload)
+            response = method(url, json=payload)
         else:
             response = method(url)
 

darwin/future/core/items/get.py

@@ -0,0 +1,64 @@
+from typing import List, Union
+from uuid import UUID
+
+from darwin.future.core.client import Client
+from darwin.future.core.types.common import QueryString
+
+
+def get_item_ids(api_client: Client, team_slug: str, dataset_id: Union[str, int]) -> List[UUID]:
+    """
+    Returns a list of item ids for the dataset
+
+    Parameters
+    ----------
+    client: Client
+        The client to use for the request
+    team_slug: str
+        The slug of the team to get item ids for
+    dataset_id: str
+        The id or slug of the dataset to get item ids for
+
+    Returns
+    -------
+    List[UUID]
+        A list of item ids
+    """
+
+    response = api_client.get(
+        f"/v2/teams/{team_slug}/items/ids",
+        QueryString({"not_statuses": "archived,error", "sort[id]": "desc", "dataset_ids": str(dataset_id)}),
+    )
+    assert type(response) == dict
+    uuids = [UUID(uuid) for uuid in response["item_ids"]]
+    return uuids
+
+
+def get_item_ids_stage(
+    api_client: Client, team_slug: str, dataset_id: Union[int, str], stage_id: Union[UUID, str]
+) -> List[UUID]:
+    """
+    Returns a list of item ids for the stage
+
+    Parameters
+    ----------
+    client: Client
+        The client to use for the request
+    team_slug: str
+        The slug of the team to get item ids for
+    dataset_id: str
+        The id or slug of the dataset to get item ids for
+    stage_id: str
+        The id or slug of the stage to get item ids for
+
+    Returns
+    -------
+    List[UUID]
+        A list of item ids
+    """
+    response = api_client.get(
+        f"/v2/teams/{team_slug}/items/ids",
+        QueryString({"workflow_stage_ids": str(stage_id), "dataset_ids": str(dataset_id)}),
+    )
+    assert type(response) == dict
+    uuids = [UUID(uuid) for uuid in response["item_ids"]]
+    return uuids

darwin/future/core/items/move_items.py

@@ -0,0 +1,41 @@
+from typing import List
+from uuid import UUID
+
+from darwin.future.core.client import Client, JSONType
+
+
+def move_items_to_stage(
+    api_client: Client, team_slug: str, workflow_id: UUID, dataset_id: int, stage_id: UUID, item_ids: List[UUID]
+) -> JSONType:
+    """
+    Moves a list of items to a stage
+
+    Parameters
+    ----------
+    client: Client
+        The client to use for the request
+    team_slug: str
+        The slug of the team to move items for
+    dataset_id: str
+        The id or slug of the dataset to move items for
+    stage_id: str
+        The id or slug of the stage to move items to
+    item_ids: List[UUID]
+        A list of item ids to move to the stage
+
+    Returns
+    -------
+    None
+    """
+
+    return api_client.post(
+        f"/v2/teams/{team_slug}/items/stage",
+        {
+            "filters": {
+                "dataset_ids": [dataset_id],
+                "item_ids": [str(id) for id in item_ids],
+            },
+            "stage_id": str(stage_id),
+            "workflow_id": str(workflow_id),
+        },
+    )

darwin/future/core/types/query.py

@@ -193,5 +193,21 @@ def where(self, *args: object, **kwargs: str) -> Query[T]:
     def collect(self) -> List[T]:
         raise NotImplementedError("Not implemented")
 
+    def collect_one(self) -> T:
+        if self.results is None:
+            self.results = list(self.collect())
+        if len(self.results) == 0:
+            raise ValueError("No results found")
+        if len(self.results) > 1:
+            raise ValueError("More than one result found")
+        return self.results[0]
+
+    def first(self) -> Optional[T]:
+        if self.results is None:
+            self.results = list(self.collect())
+        if len(self.results) == 0:
+            return None
+        return self.results[0]
+
     def _generic_execute_filter(self, objects: List[T], filter: QueryFilter) -> List[T]:
         return [m for m in objects if filter.filter_attr(getattr(m._item, filter.name))]

darwin/future/data_objects/team.py

@@ -44,7 +44,7 @@ class Team(DefaultDarwin):
     ----------
     _slug_validator: validates and auto formats the slug variable
     """
-
+    name: str
     slug: str
     id: int
     datasets: Optional[DatasetList] = None

darwin/future/data_objects/workflow.py

@@ -81,6 +81,12 @@ class WFType(Enum):
     ANNOTATE = "annotate"
     REVIEW = "review"
     COMPLETE = "complete"
+    DISCARD = "discard"
+    MODEL = "model"
+    WEBHOOK = "webhook"
+    ARCHIVE = "archive"
+    CONSENSUS_TEST = "consensus_test"
+    CONSENSUS_ENTRYPOINT = "consensus_entrypoint"
 
 
 class WFUser(DefaultDarwin):
@@ -173,7 +179,7 @@ class Workflow(DefaultDarwin):
     inserted_at: datetime
     updated_at: datetime
 
-    dataset: WFDataset
+    dataset: Optional[WFDataset]
     stages: List[WFStage]
 
     thumbnails: List[str]

darwin/future/meta/client.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 from requests.adapters import Retry
 
 from darwin.future.core.client import Client, DarwinConfig
 from darwin.future.meta.objects.team import TeamMeta
+from darwin.future.meta.objects.workflow import WorkflowMeta
+from darwin.future.meta.queries.workflow import WorkflowQuery
 
 
 class MetaClient(Client):
@@ -36,3 +38,7 @@ def team(self) -> TeamMeta:
         if self._team is None:
             self._team = TeamMeta(self)
         return self._team
+
+    # @property
+    # def workflows(self) -> WorkflowQuery:
+    #     return WorkflowQuery(self, meta_params={"team_slug": self.team.slug})

darwin/future/meta/objects/base.py

@@ -1,20 +1,22 @@
 from __future__ import annotations
 
-from typing import Generic, List, Optional, TypeVar
+from typing import Dict, Generic, List, Optional, TypeVar
 
 from darwin.future.core.client import Client
 from darwin.future.pydantic_base import DefaultDarwin
 
 R = TypeVar("R", bound=DefaultDarwin)
+Param = Dict[str, object]
 
 
 class MetaBase(Generic[R]):
     _item: Optional[R]
     client: Client
 
-    def __init__(self, client: Client, item: Optional[R] = None) -> None:
+    def __init__(self, client: Client, item: Optional[R] = None, meta_params: Optional[Param] = None) -> None:
         self.client = client
         self._item = item
+        self.meta_params = meta_params or dict()
 
     def __str__(self) -> str:
         if self._item is None:

darwin/future/meta/objects/dataset.py

@@ -1,10 +1,17 @@
-from typing import List, Optional, Tuple, Union
+from __future__ import annotations
 
+from typing import List, Optional, Sequence, Tuple, Union
+from uuid import UUID
+
+from darwin.cli_functions import upload_data
+from darwin.dataset.upload_manager import LocalFile
+from darwin.datatypes import PathLike
 from darwin.future.core.client import Client
 from darwin.future.core.datasets.create_dataset import create_dataset
 from darwin.future.core.datasets.get_dataset import get_dataset
 from darwin.future.core.datasets.list_datasets import list_datasets
 from darwin.future.core.datasets.remove_dataset import remove_dataset
+from darwin.future.core.items.get import get_item_ids
 from darwin.future.data_objects.dataset import Dataset
 from darwin.future.helpers.assertion import assert_is
 from darwin.future.meta.objects.base import MetaBase
@@ -19,6 +26,33 @@ class DatasetMeta(MetaBase[Dataset]):
     Returns:
         _type_: DatasetMeta
     """
+    @property
+    def name(self) -> str:
+        assert self._item is not None
+        assert self._item.name is not None
+        return self._item.name
+    @property
+    def slug(self) -> str:
+        assert self._item is not None
+        assert self._item.slug is not None
+        return self._item.slug
+    @property
+    def id(self) -> int:
+        assert self._item is not None
+        assert self._item.id is not None
+        return self._item.id
+
+    @property
+    def item_ids(self) -> List[UUID]:
+        """Returns a list of item ids for the dataset
+
+        Returns:
+            List[UUID]: A list of item ids
+        """
+        assert self._item is not None
+        assert self._item.id is not None
+        assert self.meta_params["team_slug"] is not None and type(self.meta_params["team_slug"]) == str
+        return get_item_ids(self.client, self.meta_params["team_slug"], str(self._item.id))
 
     def get_dataset_by_id(self) -> Dataset:
         # TODO: implement
@@ -155,3 +189,20 @@ def _validate_slug(slug: str) -> None:
 
         VALID_SLUG_CHARS = "abcdefghijklmnopqrstuvwxyz0123456789-_"
         assert_is(all(c in VALID_SLUG_CHARS for c in slug_copy), "slug must only contain valid characters")
+
+    def upload_files(
+        self,
+        files: Sequence[Union[PathLike, LocalFile]],
+        files_to_exclude: Optional[List[PathLike]] = None,
+        fps: int = 1,
+        path: Optional[str] = None,
+        frames: bool = False,
+        extract_views: bool = False,
+        preserve_folders: bool = False,
+        verbose: bool = False,
+    ) -> DatasetMeta:
+        assert self._item is not None
+        upload_data(
+            self._item.name, files, files_to_exclude, fps, path, frames, extract_views, preserve_folders, verbose
+        )
+        return self

darwin/future/meta/objects/stage.py

@@ -1,7 +1,10 @@
-from typing import Optional
+from __future__ import annotations
+
+from typing import List
 from uuid import UUID
 
-from darwin.future.core.client import Client
+from darwin.future.core.items.get import get_item_ids_stage
+from darwin.future.core.items.move_items import move_items_to_stage
 from darwin.future.data_objects.workflow import WFStage
 from darwin.future.meta.objects.base import MetaBase
 
@@ -13,11 +16,32 @@ class StageMeta(MetaBase[WFStage]):
         MetaBase (_type_): _description_
     """
 
-    def __init__(
-        self,
-        client: Client,
-        item: Optional[WFStage] = None,
-        workflow_id: Optional[UUID] = None,
-    ) -> None:
-        self._workflow_id = workflow_id
-        super().__init__(client, item)
+    @property
+    def item_ids(self) -> List[UUID]:
+        """_summary_
+
+        Returns:
+            _type_: _description_
+        """
+        assert self._item is not None
+        assert self._item.id is not None
+        return get_item_ids_stage(
+            self.client, str(self.meta_params["team_slug"]), str(self.meta_params["dataset_id"]), self.id
+        )
+
+    def move_attached_files_to_stage(self, new_stage_id: UUID) -> StageMeta:
+        assert self.meta_params["team_slug"] is not None and type(self.meta_params["team_slug"]) == str
+        assert self.meta_params["workflow_id"] is not None and type(self.meta_params["workflow_id"]) == UUID
+        assert self.meta_params["dataset_id"] is not None and type(self.meta_params["dataset_id"]) == int
+        slug, w_id, d_id = (
+            self.meta_params["team_slug"],
+            self.meta_params["workflow_id"],
+            self.meta_params["dataset_id"],
+        )
+        move_items_to_stage(self.client, slug, w_id, d_id, new_stage_id, self.item_ids)
+        return self
+
+    @property
+    def id(self) -> UUID:
+        assert self._item is not None
+        return self._item.id