arkindex-base-worker 0.3.7rc4__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

arkindex_worker/worker/element.py

@@ -2,7 +2,9 @@
 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
@@ -12,6 +14,12 @@ 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):
@@ -31,6 +39,37 @@ 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.
@@ -38,7 +77,7 @@ class ElementMixin:
         :param element_types: The missing element types to create.
         """
         for element_type in element_types:
-            self.request(
+            self.api_client.request(
                 "CreateElementType",
                 body={
                     "slug": element_type.slug,
@@ -66,10 +105,10 @@ class ElementMixin:
             isinstance(slug, str) for slug in type_slugs
         ), "Element type slugs must be strings."
 
-        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
+        if not self.corpus_types:
+            self.list_corpus_types()
 
+        missing_slugs = set(type_slugs) - set(self.corpus_types)
         if missing_slugs:
             if create_missing:
                 self.create_required_types(
@@ -79,7 +118,7 @@ class ElementMixin:
                 )
             else:
                 raise MissingTypeError(
-                    f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
+                    f'Element {pluralize("type", len(missing_slugs))} {", ".join(sorted(missing_slugs))} were not found in corpus ({self.corpus_id}).'
                 )
 
         return True
@@ -145,7 +184,7 @@ class ElementMixin:
             logger.warning("Cannot create element as this worker is in read-only mode")
             return
 
-        sub_element = self.request(
+        sub_element = self.api_client.request(
             "CreateElement",
             body={
                 "type": type,
@@ -161,10 +200,12 @@ 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.
@@ -185,6 +226,8 @@ 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):
@@ -243,14 +286,18 @@ class ElementMixin:
             logger.warning("Cannot create elements as this worker is in read-only mode")
             return
 
-        created_ids = self.request(
-            "CreateElements",
-            id=parent.id,
-            body={
-                "worker_run_id": self.worker_run_id,
-                "elements": elements,
-            },
-        )
+        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,
+                },
+            )
+        ]
 
         if self.use_cache:
             # Create the image as needed and handle both an Element and a CachedElement
@@ -311,12 +358,58 @@ class ElementMixin:
             logger.warning("Cannot link elements as this worker is in read-only mode")
             return
 
-        return self.request(
+        return self.api_client.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:
@@ -383,7 +476,7 @@ class ElementMixin:
             logger.warning("Cannot update element as this worker is in read-only mode")
             return
 
-        updated_element = self.request(
+        updated_element = self.api_client.request(
             "PartialUpdateElement",
             id=element.id,
             body=kwargs,
@@ -407,6 +500,178 @@ 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,
@@ -546,45 +811,43 @@ class ElementMixin:
                 ), "if of type bool, worker_run can only be set to False"
             query_params["worker_run"] = worker_run
 
-        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(
+        if not self.use_cache:
+            return self.api_client.paginate(
                 "ListElementChildren", id=element.id, **query_params
             )
 
-        return children
+        # 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
 
     def list_element_parents(
         self,
@@ -725,45 +988,43 @@ class ElementMixin:
                 ), "if of type bool, worker_run can only be set to False"
             query_params["worker_run"] = worker_run
 
-        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(
+        if not self.use_cache:
+            return self.api_client.paginate(
                 "ListElementParents", id=element.id, **query_params
             )
 
-        return parents
+        # 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

arkindex_worker/worker/entity.py

@@ -15,6 +15,7 @@ from arkindex_worker.cache import (
     unsupported_cache,
 )
 from arkindex_worker.models import Element, Transcription
+from arkindex_worker.utils import pluralize
 
 
 class Entity(TypedDict):
@@ -48,6 +49,7 @@ class EntityMixin:
         if not self.entity_types:
             # Load entity_types of corpus
             self.list_corpus_entity_types()
+
         for entity_type in entity_types:
             # Do nothing if type already exists
             if entity_type in self.entity_types:
@@ -60,7 +62,7 @@ class EntityMixin:
                 )
 
             # Create type if non-existent
-            self.entity_types[entity_type] = self.request(
+            self.entity_types[entity_type] = self.api_client.request(
                 "CreateEntityType",
                 body={
                     "name": entity_type,
@@ -106,7 +108,7 @@ class EntityMixin:
         entity_type_id = self.entity_types.get(type)
         assert entity_type_id, f"Entity type `{type}` not found in the corpus."
 
-        entity = self.request(
+        entity = self.api_client.request(
             "CreateEntity",
             body={
                 "name": name,
@@ -188,7 +190,7 @@ class EntityMixin:
         if confidence is not None:
             body["confidence"] = confidence
 
-        transcription_ent = self.request(
+        transcription_ent = self.api_client.request(
             "CreateTranscriptionEntity",
             id=transcription.id,
             body=body,
@@ -289,16 +291,16 @@ class EntityMixin:
             )
             return
 
-        created_ids = self.request(
+        created_entities = self.api_client.request(
             "CreateTranscriptionEntities",
             id=transcription.id,
             body={
                 "worker_run_id": self.worker_run_id,
                 "entities": entities,
             },
-        )
+        )["entities"]
 
-        return created_ids["entities"]
+        return created_entities
 
     def list_transcription_entities(
         self,
@@ -380,13 +382,12 @@ class EntityMixin:
                 "ListCorpusEntities", id=self.corpus_id, **query_params
             )
         }
+        count = len(self.entities)
         logger.info(
-            f"Loaded {len(self.entities)} entities in corpus ({self.corpus_id})"
+            f'Loaded {count} {pluralize("entity", count)} in corpus ({self.corpus_id})'
         )
 
-    def list_corpus_entity_types(
-        self,
-    ):
+    def list_corpus_entity_types(self):
         """
         Loads available entity types in corpus.
         """
@@ -396,6 +397,7 @@ class EntityMixin:
                 "ListCorpusEntityTypes", id=self.corpus_id
             )
         }
+        count = len(self.entity_types)
         logger.info(
-            f"Loaded {len(self.entity_types)} entity types in corpus ({self.corpus_id})."
+            f'Loaded {count} entity {pluralize("type", count)} in corpus ({self.corpus_id}).'
         )

arkindex_worker/worker/image.py

@@ -0,0 +1,21 @@
+"""
+ElementsWorker methods for images.
+"""
+
+from arkindex_worker.models import Image
+
+
+class ImageMixin:
+    def create_iiif_url(self, url: str) -> Image:
+        """
+        Create an image from an existing IIIF image by URL.
+        The URL should be of the image's identifier, not of its Image Information request (`/info.json`).
+
+        :param url: URL of the image.
+        :returns: The created image.
+        """
+        assert url and isinstance(
+            url, str
+        ), "url shouldn't be null and should be of type str"
+
+        return Image(self.api_client.request("CreateIIIFURL", body={"url": url}))