arkindex-base-worker 0.4.0__py3-none-any.whl → 0.4.0a2__py3-none-any.whl

arkindex_worker/worker/dataset.py

@@ -2,8 +2,6 @@
 BaseWorker methods for datasets.
 """
 
-import uuid
-from argparse import ArgumentTypeError
 from collections.abc import Iterator
 from enum import Enum
 
@@ -38,55 +36,7 @@ class DatasetState(Enum):
     """
 
 
-class MissingDatasetArchive(Exception):
-    """
-    Exception raised when the compressed archive associated to
-    a dataset isn't found in its task artifacts.
-    """
-
-
-def check_dataset_set(value: str) -> tuple[uuid.UUID, str]:
-    """The `--set` argument should have the following format:
-    <dataset_id>:<set_name>
-
-    Args:
-        value (str): Provided argument.
-
-    Raises:
-        ArgumentTypeError: When the value is invalid.
-
-    Returns:
-        tuple[uuid.UUID, str]: The ID of the dataset parsed as UUID and the name of the set.
-    """
-    values = value.split(":")
-    if len(values) != 2:
-        raise ArgumentTypeError(
-            f"'{value}' is not in the correct format `<dataset_id>:<set_name>`"
-        )
-
-    dataset_id, set_name = values
-    try:
-        dataset_id = uuid.UUID(dataset_id)
-        return (dataset_id, set_name)
-    except (TypeError, ValueError) as e:
-        raise ArgumentTypeError(f"'{dataset_id}' should be a valid UUID") from e
-
-
 class DatasetMixin:
-    def add_arguments(self) -> None:
-        """Define specific ``argparse`` arguments for the worker using this mixin"""
-        self.parser.add_argument(
-            "--set",
-            type=check_dataset_set,
-            nargs="+",
-            help="""
-                One or more Arkindex dataset sets, format is <dataset_uuid>:<set_name>
-                (e.g.: "12341234-1234-1234-1234-123412341234:train")
-            """,
-            default=[],
-        )
-        super().add_arguments()
-
     def list_process_sets(self) -> Iterator[Set]:
         """
         List dataset sets associated to the worker's process. This helper is not available in developer mode.
@@ -123,26 +73,6 @@ class DatasetMixin:
 
         return map(lambda result: Element(**result["element"]), results)
 
-    def list_sets(self) -> Iterator[Set]:
-        """
-        List the sets to be processed, either from the CLI arguments or using the
-        [list_process_sets][arkindex_worker.worker.dataset.DatasetMixin.list_process_sets] method.
-
-        :returns: An iterator of ``Set`` objects.
-        """
-        if not self.is_read_only:
-            yield from self.list_process_sets()
-
-        datasets: dict[uuid.UUID, Dataset] = {}
-        for dataset_id, set_name in self.args.set:
-            # Retrieving dataset information if not already cached
-            if dataset_id not in datasets:
-                datasets[dataset_id] = Dataset(
-                    **self.api_client.request("RetrieveDataset", id=dataset_id)
-                )
-
-            yield Set(name=set_name, dataset=datasets[dataset_id])
-
     @unsupported_cache
     def update_dataset_state(self, dataset: Dataset, state: DatasetState) -> Dataset:
         """
@@ -163,7 +93,7 @@ class DatasetMixin:
             logger.warning("Cannot update dataset as this worker is in read-only mode")
             return
 
-        updated_dataset = self.api_client.request(
+        updated_dataset = self.request(
             "PartialUpdateDataset",
             id=dataset.id,
             body={"state": state.value},

arkindex_worker/worker/element.py

@@ -2,9 +2,7 @@
 ElementsWorker methods for elements and element types.
 """
 
-import os
 from collections.abc import Iterable
-from operator import attrgetter
 from typing import NamedTuple
 from uuid import UUID
 from warnings import warn
@@ -14,12 +12,6 @@ from peewee import IntegrityError
 from arkindex_worker import logger
 from arkindex_worker.cache import CachedElement, CachedImage, unsupported_cache
 from arkindex_worker.models import Element
-from arkindex_worker.utils import (
-    DEFAULT_BATCH_SIZE,
-    batch_publication,
-    make_batches,
-    pluralize,
-)
 
 
 class ElementType(NamedTuple):
@@ -39,37 +31,6 @@ class MissingTypeError(Exception):
 
 
 class ElementMixin:
-    def add_arguments(self):
-        """Define specific ``argparse`` arguments for the worker using this mixin"""
-        self.parser.add_argument(
-            "--elements-list",
-            help="JSON elements list to use",
-            type=open,
-            default=os.environ.get("TASK_ELEMENTS"),
-        )
-        self.parser.add_argument(
-            "--element",
-            type=str,
-            nargs="+",
-            help="One or more Arkindex element ID",
-        )
-        super().add_arguments()
-
-    def list_corpus_types(self):
-        """
-        Loads available element types in corpus.
-        """
-        self.corpus_types = {
-            element_type["slug"]: element_type
-            for element_type in self.api_client.request(
-                "RetrieveCorpus", id=self.corpus_id
-            )["types"]
-        }
-        count = len(self.corpus_types)
-        logger.info(
-            f'Loaded {count} element {pluralize("type", count)} in corpus ({self.corpus_id}).'
-        )
-
     @unsupported_cache
     def create_required_types(self, element_types: list[ElementType]):
         """Creates given element types in the corpus.
@@ -77,7 +38,7 @@ class ElementMixin:
         :param element_types: The missing element types to create.
         """
         for element_type in element_types:
-            self.api_client.request(
+            self.request(
                 "CreateElementType",
                 body={
                     "slug": element_type.slug,
@@ -105,10 +66,10 @@ class ElementMixin:
             isinstance(slug, str) for slug in type_slugs
         ), "Element type slugs must be strings."
 
-        if not self.corpus_types:
-            self.list_corpus_types()
+        corpus = self.request("RetrieveCorpus", id=self.corpus_id)
+        available_slugs = {element_type["slug"] for element_type in corpus["types"]}
+        missing_slugs = set(type_slugs) - available_slugs
 
-        missing_slugs = set(type_slugs) - set(self.corpus_types)
         if missing_slugs:
             if create_missing:
                 self.create_required_types(
@@ -118,7 +79,7 @@ class ElementMixin:
                 )
             else:
                 raise MissingTypeError(
-                    f'Element {pluralize("type", len(missing_slugs))} {", ".join(sorted(missing_slugs))} were not found in corpus ({self.corpus_id}).'
+                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
                 )
 
         return True
@@ -184,7 +145,7 @@ class ElementMixin:
             logger.warning("Cannot create element as this worker is in read-only mode")
             return
 
-        sub_element = self.api_client.request(
+        sub_element = self.request(
             "CreateElement",
             body={
                 "type": type,
@@ -200,12 +161,10 @@ class ElementMixin:
 
         return sub_element["id"] if slim_output else sub_element
 
-    @batch_publication
     def create_elements(
         self,
         parent: Element | CachedElement,
         elements: list[dict[str, str | list[list[int | float]] | float | None]],
-        batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> list[dict[str, str]]:
         """
         Create child elements on the given element in a single API request.
@@ -226,8 +185,6 @@ class ElementMixin:
             confidence (float or None)
                 Optional confidence score, between 0.0 and 1.0.
 
-        :param batch_size: The size of each batch, which will be used to split the publication to avoid API errors.
-
         :return: List of dicts, with each dict having a single key, ``id``, holding the UUID of each created element.
         """
         if isinstance(parent, Element):
@@ -286,18 +243,14 @@ class ElementMixin:
             logger.warning("Cannot create elements as this worker is in read-only mode")
             return
 
-        created_ids = [
-            created_id
-            for batch in make_batches(elements, "element", batch_size)
-            for created_id in self.api_client.request(
-                "CreateElements",
-                id=parent.id,
-                body={
-                    "worker_run_id": self.worker_run_id,
-                    "elements": batch,
-                },
-            )
-        ]
+        created_ids = self.request(
+            "CreateElements",
+            id=parent.id,
+            body={
+                "worker_run_id": self.worker_run_id,
+                "elements": elements,
+            },
+        )
 
         if self.use_cache:
             # Create the image as needed and handle both an Element and a CachedElement
@@ -358,58 +311,12 @@ class ElementMixin:
             logger.warning("Cannot link elements as this worker is in read-only mode")
             return
 
-        return self.api_client.request(
+        return self.request(
             "CreateElementParent",
             parent=parent.id,
             child=child.id,
         )
 
-    @unsupported_cache
-    @batch_publication
-    def create_element_children(
-        self,
-        parent: Element,
-        children: list[Element],
-        batch_size: int = DEFAULT_BATCH_SIZE,
-    ) -> list[str]:
-        """
-        Link multiple elements to a single parent through the API.
-
-        :param parent: Parent element.
-        :param children: A list of child elements.
-        :param batch_size: The size of each batch, which will be used to split the publication to avoid API errors.
-
-        :returns: A list containing the string UUID of each child linked to the parent.
-        """
-        assert parent and isinstance(
-            parent, Element
-        ), "parent shouldn't be null and should be of type Element"
-
-        assert children and isinstance(
-            children, list
-        ), "children shouldn't be null and should be of type list"
-
-        for index, child in enumerate(children):
-            assert isinstance(
-                child, Element
-            ), f"Child at index {index} in children: Should be of type Element"
-
-        if self.is_read_only:
-            logger.warning("Cannot link elements as this worker is in read-only mode")
-            return
-
-        return [
-            child_id
-            for batch in make_batches(children, "child", batch_size)
-            for child_id in self.api_client.request(
-                "CreateElementChildren",
-                id=parent.id,
-                body={
-                    "children": list(map(attrgetter("id"), batch)),
-                },
-            )["children"]
-        ]
-
     def partial_update_element(
         self, element: Element | CachedElement, **kwargs
     ) -> dict:
@@ -476,7 +383,7 @@ class ElementMixin:
             logger.warning("Cannot update element as this worker is in read-only mode")
             return
 
-        updated_element = self.api_client.request(
+        updated_element = self.request(
             "PartialUpdateElement",
             id=element.id,
             body=kwargs,
@@ -500,178 +407,6 @@ class ElementMixin:
 
         return updated_element
 
-    def list_elements(
-        self,
-        folder: bool | None = None,
-        name: str | None = None,
-        top_level: bool | None = None,
-        transcription_worker_version: str | bool | None = None,
-        transcription_worker_run: str | bool | None = None,
-        type: str | None = None,
-        with_classes: bool | None = None,
-        with_corpus: bool | None = None,
-        with_metadata: bool | None = None,
-        with_has_children: bool | None = None,
-        with_zone: bool | None = None,
-        worker_version: str | bool | None = None,
-        worker_run: str | bool | None = None,
-    ) -> Iterable[dict] | Iterable[CachedElement]:
-        """
-        List element in a corpus.
-
-        Warns:
-        ----
-        The following parameters are **deprecated**:
-
-        - `transcription_worker_version` in favor of `transcription_worker_run`
-        - `worker_version` in favor of `worker_run`
-
-        :param folder: Restrict to or exclude elements with folder types.
-           This parameter is not supported when caching is enabled.
-        :param name: Restrict to elements whose name contain a substring (case-insensitive).
-           This parameter is not supported when caching is enabled.
-        :param top_level: Restrict to or exclude folder elements without parent elements (top-level elements).
-           This parameter is not supported when caching is enabled.
-        :param transcription_worker_version: **Deprecated** Restrict to elements that have a transcription created by a worker version with this UUID. Set to False to look for elements that have a manual transcription.
-           This parameter is not supported when caching is enabled.
-        :param transcription_worker_run: Restrict to elements that have a transcription created by a worker run with this UUID. Set to False to look for elements that have a manual transcription.
-           This parameter is not supported when caching is enabled.
-        :param type: Restrict to elements with a specific type slug
-           This parameter is not supported when caching is enabled.
-        :param with_classes: Include each element's classifications in the response.
-           This parameter is not supported when caching is enabled.
-        :param with_corpus: Include each element's corpus in the response.
-           This parameter is not supported when caching is enabled.
-        :param with_has_children: Include the ``has_children`` attribute in the response,
-           indicating if this element has child elements of its own.
-           This parameter is not supported when caching is enabled.
-        :param with_metadata: Include each element's metadata in the response.
-           This parameter is not supported when caching is enabled.
-        :param with_zone: Include the ``zone`` attribute in the response,
-           holding the element's image and polygon.
-           This parameter is not supported when caching is enabled.
-        :param worker_version: **Deprecated** Restrict to elements created by a worker version with this UUID.
-        :param worker_run: Restrict to elements created by a worker run with this UUID.
-        :return: An iterable of dicts from the ``ListElementChildren`` API endpoint,
-           or an iterable of [CachedElement][arkindex_worker.cache.CachedElement] when caching is enabled.
-        """
-        query_params = {}
-        if folder is not None:
-            assert isinstance(folder, bool), "folder should be of type bool"
-            query_params["folder"] = folder
-        if name:
-            assert isinstance(name, str), "name should be of type str"
-            query_params["name"] = name
-        if top_level is not None:
-            assert isinstance(top_level, bool), "top_level should be of type bool"
-            query_params["top_level"] = top_level
-        if transcription_worker_version is not None:
-            warn(
-                "`transcription_worker_version` usage is deprecated. Consider using `transcription_worker_run` instead.",
-                DeprecationWarning,
-                stacklevel=1,
-            )
-            assert isinstance(
-                transcription_worker_version, str | bool
-            ), "transcription_worker_version should be of type str or bool"
-            if isinstance(transcription_worker_version, bool):
-                assert (
-                    transcription_worker_version is False
-                ), "if of type bool, transcription_worker_version can only be set to False"
-            query_params["transcription_worker_version"] = transcription_worker_version
-        if transcription_worker_run is not None:
-            assert isinstance(
-                transcription_worker_run, str | bool
-            ), "transcription_worker_run should be of type str or bool"
-            if isinstance(transcription_worker_run, bool):
-                assert (
-                    transcription_worker_run is False
-                ), "if of type bool, transcription_worker_run can only be set to False"
-            query_params["transcription_worker_run"] = transcription_worker_run
-        if type:
-            assert isinstance(type, str), "type should be of type str"
-            query_params["type"] = type
-        if with_classes is not None:
-            assert isinstance(with_classes, bool), "with_classes should be of type bool"
-            query_params["with_classes"] = with_classes
-        if with_corpus is not None:
-            assert isinstance(with_corpus, bool), "with_corpus should be of type bool"
-            query_params["with_corpus"] = with_corpus
-        if with_has_children is not None:
-            assert isinstance(
-                with_has_children, bool
-            ), "with_has_children should be of type bool"
-            query_params["with_has_children"] = with_has_children
-        if with_metadata is not None:
-            assert isinstance(
-                with_metadata, bool
-            ), "with_metadata should be of type bool"
-            query_params["with_metadata"] = with_metadata
-        if with_zone is not None:
-            assert isinstance(with_zone, bool), "with_zone should be of type bool"
-            query_params["with_zone"] = with_zone
-        if worker_version is not None:
-            warn(
-                "`worker_version` usage is deprecated. Consider using `worker_run` instead.",
-                DeprecationWarning,
-                stacklevel=1,
-            )
-            assert isinstance(
-                worker_version, str | bool
-            ), "worker_version should be of type str or bool"
-            if isinstance(worker_version, bool):
-                assert (
-                    worker_version is False
-                ), "if of type bool, worker_version can only be set to False"
-            query_params["worker_version"] = worker_version
-        if worker_run is not None:
-            assert isinstance(
-                worker_run, str | bool
-            ), "worker_run should be of type str or bool"
-            if isinstance(worker_run, bool):
-                assert (
-                    worker_run is False
-                ), "if of type bool, worker_run can only be set to False"
-            query_params["worker_run"] = worker_run
-
-        if not self.use_cache:
-            return self.api_client.paginate(
-                "ListElements", corpus=self.corpus_id, **query_params
-            )
-
-        # Checking that we only received query_params handled by the cache
-        assert (
-            set(query_params.keys())
-            <= {
-                "type",
-                "worker_version",
-                "worker_run",
-            }
-        ), "When using the local cache, you can only filter by 'type' and/or 'worker_version' and/or 'worker_run'"
-
-        query = CachedElement.select()
-        if type:
-            query = query.where(CachedElement.type == type)
-        if worker_version is not None:
-            # If worker_version=False, filter by manual worker_version e.g. None
-            worker_version_id = worker_version or None
-            if worker_version_id:
-                query = query.where(
-                    CachedElement.worker_version_id == worker_version_id
-                )
-            else:
-                query = query.where(CachedElement.worker_version_id.is_null())
-
-        if worker_run is not None:
-            # If worker_run=False, filter by manual worker_run e.g. None
-            worker_run_id = worker_run or None
-            if worker_run_id:
-                query = query.where(CachedElement.worker_run_id == worker_run_id)
-            else:
-                query = query.where(CachedElement.worker_run_id.is_null())
-
-        return query
-
     def list_element_children(
         self,
         element: Element | CachedElement,
@@ -811,43 +546,45 @@ class ElementMixin:
                 ), "if of type bool, worker_run can only be set to False"
             query_params["worker_run"] = worker_run
 
-        if not self.use_cache:
-            return self.api_client.paginate(
+        if self.use_cache:
+            # Checking that we only received query_params handled by the cache
+            assert (
+                set(query_params.keys())
+                <= {
+                    "type",
+                    "worker_version",
+                    "worker_run",
+                }
+            ), "When using the local cache, you can only filter by 'type' and/or 'worker_version' and/or 'worker_run'"
+
+            query = CachedElement.select().where(CachedElement.parent_id == element.id)
+            if type:
+                query = query.where(CachedElement.type == type)
+            if worker_version is not None:
+                # If worker_version=False, filter by manual worker_version e.g. None
+                worker_version_id = worker_version or None
+                if worker_version_id:
+                    query = query.where(
+                        CachedElement.worker_version_id == worker_version_id
+                    )
+                else:
+                    query = query.where(CachedElement.worker_version_id.is_null())
+
+            if worker_run is not None:
+                # If worker_run=False, filter by manual worker_run e.g. None
+                worker_run_id = worker_run or None
+                if worker_run_id:
+                    query = query.where(CachedElement.worker_run_id == worker_run_id)
+                else:
+                    query = query.where(CachedElement.worker_run_id.is_null())
+
+            return query
+        else:
+            children = self.api_client.paginate(
                 "ListElementChildren", id=element.id, **query_params
             )
 
-        # Checking that we only received query_params handled by the cache
-        assert (
-            set(query_params.keys())
-            <= {
-                "type",
-                "worker_version",
-                "worker_run",
-            }
-        ), "When using the local cache, you can only filter by 'type' and/or 'worker_version' and/or 'worker_run'"
-
-        query = CachedElement.select().where(CachedElement.parent_id == element.id)
-        if type:
-            query = query.where(CachedElement.type == type)
-        if worker_version is not None:
-            # If worker_version=False, filter by manual worker_version e.g. None
-            worker_version_id = worker_version or None
-            if worker_version_id:
-                query = query.where(
-                    CachedElement.worker_version_id == worker_version_id
-                )
-            else:
-                query = query.where(CachedElement.worker_version_id.is_null())
-
-        if worker_run is not None:
-            # If worker_run=False, filter by manual worker_run e.g. None
-            worker_run_id = worker_run or None
-            if worker_run_id:
-                query = query.where(CachedElement.worker_run_id == worker_run_id)
-            else:
-                query = query.where(CachedElement.worker_run_id.is_null())
-
-        return query
+        return children
 
     def list_element_parents(
         self,
@@ -988,43 +725,45 @@ class ElementMixin:
                 ), "if of type bool, worker_run can only be set to False"
             query_params["worker_run"] = worker_run
 
-        if not self.use_cache:
-            return self.api_client.paginate(
+        if self.use_cache:
+            # Checking that we only received query_params handled by the cache
+            assert (
+                set(query_params.keys())
+                <= {
+                    "type",
+                    "worker_version",
+                    "worker_run",
+                }
+            ), "When using the local cache, you can only filter by 'type' and/or 'worker_version' and/or 'worker_run'"
+
+            parent_ids = CachedElement.select(CachedElement.parent_id).where(
+                CachedElement.id == element.id
+            )
+            query = CachedElement.select().where(CachedElement.id.in_(parent_ids))
+            if type:
+                query = query.where(CachedElement.type == type)
+            if worker_version is not None:
+                # If worker_version=False, filter by manual worker_version e.g. None
+                worker_version_id = worker_version or None
+                if worker_version_id:
+                    query = query.where(
+                        CachedElement.worker_version_id == worker_version_id
+                    )
+                else:
+                    query = query.where(CachedElement.worker_version_id.is_null())
+
+            if worker_run is not None:
+                # If worker_run=False, filter by manual worker_run e.g. None
+                worker_run_id = worker_run or None
+                if worker_run_id:
+                    query = query.where(CachedElement.worker_run_id == worker_run_id)
+                else:
+                    query = query.where(CachedElement.worker_run_id.is_null())
+
+            return query
+        else:
+            parents = self.api_client.paginate(
                 "ListElementParents", id=element.id, **query_params
             )
 
-        # Checking that we only received query_params handled by the cache
-        assert (
-            set(query_params.keys())
-            <= {
-                "type",
-                "worker_version",
-                "worker_run",
-            }
-        ), "When using the local cache, you can only filter by 'type' and/or 'worker_version' and/or 'worker_run'"
-
-        parent_ids = CachedElement.select(CachedElement.parent_id).where(
-            CachedElement.id == element.id
-        )
-        query = CachedElement.select().where(CachedElement.id.in_(parent_ids))
-        if type:
-            query = query.where(CachedElement.type == type)
-        if worker_version is not None:
-            # If worker_version=False, filter by manual worker_version e.g. None
-            worker_version_id = worker_version or None
-            if worker_version_id:
-                query = query.where(
-                    CachedElement.worker_version_id == worker_version_id
-                )
-            else:
-                query = query.where(CachedElement.worker_version_id.is_null())
-
-        if worker_run is not None:
-            # If worker_run=False, filter by manual worker_run e.g. None
-            worker_run_id = worker_run or None
-            if worker_run_id:
-                query = query.where(CachedElement.worker_run_id == worker_run_id)
-            else:
-                query = query.where(CachedElement.worker_run_id.is_null())
-
-        return query
+        return parents