arkindex-base-worker 0.3.6rc1__py3-none-any.whl → 0.3.6rc2__py3-none-any.whl

arkindex_worker/worker/base.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 The base class for all Arkindex workers.
 """
@@ -9,7 +8,6 @@ import os
 import shutil
 from pathlib import Path
 from tempfile import mkdtemp
-from typing import List, Optional
 
 import gnupg
 import yaml
@@ -52,15 +50,15 @@ class ExtrasDirNotFoundError(Exception):
     """
 
 
-class BaseWorker(object):
+class BaseWorker:
     """
     Base class for Arkindex workers.
     """
 
     def __init__(
         self,
-        description: Optional[str] = "Arkindex Base Worker",
-        support_cache: Optional[bool] = False,
+        description: str | None = "Arkindex Base Worker",
+        support_cache: bool | None = False,
     ):
         """
         Initialize the worker.
@@ -217,6 +215,9 @@ class BaseWorker(object):
         # Define model_version_id from environment
         self.model_version_id = os.environ.get("ARKINDEX_MODEL_VERSION_ID")
 
+        # Define model_details from environment
+        self.model_details = {"id": os.environ.get("ARKINDEX_MODEL_ID")}
+
         # Load all required secrets
         self.secrets = {name: self.load_secret(Path(name)) for name in required_secrets}
 
@@ -259,6 +260,9 @@ class BaseWorker(object):
             # Set model_version ID as worker attribute
             self.model_version_id = model_version.get("id")
 
+            # Set model details as worker attribute
+            self.model_details = model_version.get("model")
+
         # Retrieve initial configuration from API
         self.config = worker_version["configuration"].get("configuration", {})
         if "user_configuration" in worker_version["configuration"]:
@@ -347,7 +351,8 @@ class BaseWorker(object):
 
             try:
                 gpg = gnupg.GPG()
-                decrypted = gpg.decrypt_file(open(path, "rb"))
+                with path.open("rb") as gpg_file:
+                    decrypted = gpg.decrypt_file(gpg_file)
                 assert (
                     decrypted.ok
                 ), f"GPG error: {decrypted.status} - {decrypted.stderr}"
@@ -406,7 +411,7 @@ class BaseWorker(object):
                 )
             return extras_dir
 
-    def find_parents_file_paths(self, filename: Path) -> List[Path]:
+    def find_parents_file_paths(self, filename: Path) -> list[Path]:
         """
         Find the paths of a specific file from the parent tasks.
         Only works if the task_parents attributes is updated, so if the cache is supported,

arkindex_worker/worker/classification.py

@@ -1,9 +1,7 @@
-# -*- coding: utf-8 -*-
 """
 ElementsWorker methods for classifications and ML classes.
 """
 
-from typing import Dict, List, Optional, Union
 from uuid import UUID
 
 from apistar.exceptions import ErrorResponse
@@ -14,7 +12,7 @@ from arkindex_worker.cache import CachedClassification, CachedElement
 from arkindex_worker.models import Element
 
 
-class ClassificationMixin(object):
+class ClassificationMixin:
     def load_corpus_classes(self):
         """
         Load all ML classes available in the worker's corpus and store them in the ``self.classes`` cache.
@@ -91,11 +89,11 @@ class ClassificationMixin(object):
 
     def create_classification(
         self,
-        element: Union[Element, CachedElement],
+        element: Element | CachedElement,
         ml_class: str,
         confidence: float,
-        high_confidence: Optional[bool] = False,
-    ) -> Dict[str, str]:
+        high_confidence: bool = False,
+    ) -> dict[str, str]:
         """
         Create a classification on the given element through the API.
 
@@ -106,7 +104,7 @@ class ClassificationMixin(object):
         :returns: The created classification, as returned by the ``CreateClassification`` API endpoint.
         """
         assert element and isinstance(
-            element, (Element, CachedElement)
+            element, Element | CachedElement
         ), "element shouldn't be null and should be an Element or CachedElement"
         assert ml_class and isinstance(
             ml_class, str
@@ -180,9 +178,9 @@ class ClassificationMixin(object):
 
     def create_classifications(
         self,
-        element: Union[Element, CachedElement],
-        classifications: List[Dict[str, Union[str, float, bool]]],
-    ) -> List[Dict[str, Union[str, float, bool]]]:
+        element: Element | CachedElement,
+        classifications: list[dict[str, str | float | bool]],
+    ) -> list[dict[str, str | float | bool]]:
         """
         Create multiple classifications at once on the given element through the API.
 
@@ -196,7 +194,7 @@ class ClassificationMixin(object):
            the ``CreateClassifications`` API endpoint.
         """
         assert element and isinstance(
-            element, (Element, CachedElement)
+            element, Element | CachedElement
         ), "element shouldn't be null and should be an Element or CachedElement"
         assert classifications and isinstance(
             classifications, list
@@ -204,17 +202,17 @@ class ClassificationMixin(object):
 
         for index, classification in enumerate(classifications):
             ml_class_id = classification.get("ml_class_id")
-            assert ml_class_id and isinstance(
-                ml_class_id, str
+            assert (
+                ml_class_id and isinstance(ml_class_id, str)
             ), f"Classification at index {index} in classifications: ml_class_id shouldn't be null and should be of type str"
 
             # Make sure it's a valid UUID
             try:
                 UUID(ml_class_id)
-            except ValueError:
+            except ValueError as e:
                 raise ValueError(
                     f"Classification at index {index} in classifications: ml_class_id is not a valid uuid."
-                )
+                ) from e
 
             confidence = classification.get("confidence")
             assert (

arkindex_worker/worker/dataset.py

@@ -1,10 +1,9 @@
-# -*- coding: utf-8 -*-
 """
 BaseWorker methods for datasets.
 """
 
+from collections.abc import Iterator
 from enum import Enum
-from typing import Iterator, Tuple
 
 from arkindex_worker import logger
 from arkindex_worker.models import Dataset, Element
@@ -36,7 +35,7 @@ class DatasetState(Enum):
     """
 
 
-class DatasetMixin(object):
+class DatasetMixin:
     def list_process_datasets(self) -> Iterator[Dataset]:
         """
         List datasets associated to the worker's process. This helper is not available in developer mode.
@@ -51,7 +50,7 @@ class DatasetMixin(object):
 
         return map(Dataset, list(results))
 
-    def list_dataset_elements(self, dataset: Dataset) -> Iterator[Tuple[str, Element]]:
+    def list_dataset_elements(self, dataset: Dataset) -> Iterator[tuple[str, Element]]:
         """
         List elements in a dataset.
 

arkindex_worker/worker/element.py

@@ -1,8 +1,8 @@
-# -*- coding: utf-8 -*-
 """
 ElementsWorker methods for elements and element types.
 """
-from typing import Dict, Iterable, List, NamedTuple, Optional, Union
+from collections.abc import Iterable
+from typing import NamedTuple
 from uuid import UUID
 
 from peewee import IntegrityError
@@ -28,8 +28,8 @@ class MissingTypeError(Exception):
     """
 
 
-class ElementMixin(object):
-    def create_required_types(self, element_types: List[ElementType]):
+class ElementMixin:
+    def create_required_types(self, element_types: list[ElementType]):
         """Creates given element types in the corpus.
 
         :param element_types: The missing element types to create.
@@ -86,9 +86,9 @@ class ElementMixin(object):
         element: Element,
         type: str,
         name: str,
-        polygon: List[List[Union[int, float]]],
-        confidence: Optional[float] = None,
-        slim_output: Optional[bool] = True,
+        polygon: list[list[int | float]],
+        confidence: float | None = None,
+        slim_output: bool = True,
     ) -> str:
         """
         Create a child element on the given element through the API.
@@ -117,7 +117,7 @@ class ElementMixin(object):
             isinstance(point, list) and len(point) == 2 for point in polygon
         ), "polygon points should be lists of two items"
         assert all(
-            isinstance(coord, (int, float)) for point in polygon for coord in point
+            isinstance(coord, int | float) for point in polygon for coord in point
         ), "polygon points should be lists of two numbers"
         assert confidence is None or (
             isinstance(confidence, float) and 0 <= confidence <= 1
@@ -146,11 +146,9 @@ class ElementMixin(object):
 
     def create_elements(
         self,
-        parent: Union[Element, CachedElement],
-        elements: List[
-            Dict[str, Union[str, List[List[Union[int, float]]], float, None]]
-        ],
-    ) -> List[Dict[str, str]]:
+        parent: Element | CachedElement,
+        elements: list[dict[str, str | list[list[int | float]] | float | None]],
+    ) -> list[dict[str, str]]:
         """
         Create child elements on the given element in a single API request.
 
@@ -195,18 +193,18 @@ class ElementMixin(object):
             ), f"Element at index {index} in elements: Should be of type dict"
 
             name = element.get("name")
-            assert name and isinstance(
-                name, str
+            assert (
+                name and isinstance(name, str)
             ), f"Element at index {index} in elements: name shouldn't be null and should be of type str"
 
             type = element.get("type")
-            assert type and isinstance(
-                type, str
+            assert (
+                type and isinstance(type, str)
             ), f"Element at index {index} in elements: type shouldn't be null and should be of type str"
 
             polygon = element.get("polygon")
-            assert polygon and isinstance(
-                polygon, list
+            assert (
+                polygon and isinstance(polygon, list)
             ), f"Element at index {index} in elements: polygon shouldn't be null and should be of type list"
             assert (
                 len(polygon) >= 3
@@ -215,12 +213,13 @@ class ElementMixin(object):
                 isinstance(point, list) and len(point) == 2 for point in polygon
             ), f"Element at index {index} in elements: polygon points should be lists of two items"
             assert all(
-                isinstance(coord, (int, float)) for point in polygon for coord in point
+                isinstance(coord, int | float) for point in polygon for coord in point
             ), f"Element at index {index} in elements: polygon points should be lists of two numbers"
 
             confidence = element.get("confidence")
-            assert confidence is None or (
-                isinstance(confidence, float) and 0 <= confidence <= 1
+            assert (
+                confidence is None
+                or (isinstance(confidence, float) and 0 <= confidence <= 1)
             ), f"Element at index {index} in elements: confidence should be None or a float in [0..1] range"
 
         if self.is_read_only:
@@ -272,7 +271,7 @@ class ElementMixin(object):
         return created_ids
 
     def partial_update_element(
-        self, element: Union[Element, CachedElement], **kwargs
+        self, element: Element | CachedElement, **kwargs
     ) -> dict:
         """
         Partially updates an element through the API.
@@ -292,7 +291,7 @@ class ElementMixin(object):
         :returns: A dict from the ``PartialUpdateElement`` API endpoint,
         """
         assert element and isinstance(
-            element, (Element, CachedElement)
+            element, Element | CachedElement
         ), "element shouldn't be null and should be an Element or CachedElement"
 
         if "type" in kwargs:
@@ -309,7 +308,7 @@ class ElementMixin(object):
                 isinstance(point, list) and len(point) == 2 for point in polygon
             ), "polygon points should be lists of two items"
             assert all(
-                isinstance(coord, (int, float)) for point in polygon for coord in point
+                isinstance(coord, int | float) for point in polygon for coord in point
             ), "polygon points should be lists of two numbers"
 
         if "confidence" in kwargs:
@@ -363,21 +362,21 @@ class ElementMixin(object):
 
     def list_element_children(
         self,
-        element: Union[Element, CachedElement],
-        folder: Optional[bool] = None,
-        name: Optional[str] = None,
-        recursive: Optional[bool] = None,
-        transcription_worker_version: Optional[Union[str, bool]] = None,
-        transcription_worker_run: Optional[Union[str, bool]] = None,
-        type: Optional[str] = None,
-        with_classes: Optional[bool] = None,
-        with_corpus: Optional[bool] = None,
-        with_metadata: Optional[bool] = None,
-        with_has_children: Optional[bool] = None,
-        with_zone: Optional[bool] = None,
-        worker_version: Optional[Union[str, bool]] = None,
-        worker_run: Optional[Union[str, bool]] = None,
-    ) -> Union[Iterable[dict], Iterable[CachedElement]]:
+        element: Element | CachedElement,
+        folder: bool | None = None,
+        name: str | None = None,
+        recursive: 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 children of an element.
 
@@ -412,7 +411,7 @@ class ElementMixin(object):
            or an iterable of [CachedElement][arkindex_worker.cache.CachedElement] when caching is enabled.
         """
         assert element and isinstance(
-            element, (Element, CachedElement)
+            element, Element | CachedElement
         ), "element shouldn't be null and should be an Element or CachedElement"
         query_params = {}
         if folder is not None:
@@ -426,7 +425,7 @@ class ElementMixin(object):
             query_params["recursive"] = recursive
         if transcription_worker_version is not None:
             assert isinstance(
-                transcription_worker_version, (str, bool)
+                transcription_worker_version, str | bool
             ), "transcription_worker_version should be of type str or bool"
             if isinstance(transcription_worker_version, bool):
                 assert (
@@ -435,7 +434,7 @@ class ElementMixin(object):
             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, str | bool
             ), "transcription_worker_run should be of type str or bool"
             if isinstance(transcription_worker_run, bool):
                 assert (
@@ -466,7 +465,7 @@ class ElementMixin(object):
             query_params["with_zone"] = with_zone
         if worker_version is not None:
             assert isinstance(
-                worker_version, (str, bool)
+                worker_version, str | bool
             ), "worker_version should be of type str or bool"
             if isinstance(worker_version, bool):
                 assert (
@@ -475,7 +474,7 @@ class ElementMixin(object):
             query_params["worker_version"] = worker_version
         if worker_run is not None:
             assert isinstance(
-                worker_run, (str, bool)
+                worker_run, str | bool
             ), "worker_run should be of type str or bool"
             if isinstance(worker_run, bool):
                 assert (
@@ -485,11 +484,14 @@ class ElementMixin(object):
 
         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'"
+            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:
@@ -522,21 +524,21 @@ class ElementMixin(object):
 
     def list_element_parents(
         self,
-        element: Union[Element, CachedElement],
-        folder: Optional[bool] = None,
-        name: Optional[str] = None,
-        recursive: Optional[bool] = None,
-        transcription_worker_version: Optional[Union[str, bool]] = None,
-        transcription_worker_run: Optional[Union[str, bool]] = None,
-        type: Optional[str] = None,
-        with_classes: Optional[bool] = None,
-        with_corpus: Optional[bool] = None,
-        with_metadata: Optional[bool] = None,
-        with_has_children: Optional[bool] = None,
-        with_zone: Optional[bool] = None,
-        worker_version: Optional[Union[str, bool]] = None,
-        worker_run: Optional[Union[str, bool]] = None,
-    ) -> Union[Iterable[dict], Iterable[CachedElement]]:
+        element: Element | CachedElement,
+        folder: bool | None = None,
+        name: str | None = None,
+        recursive: 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 parents of an element.
 
@@ -571,7 +573,7 @@ class ElementMixin(object):
            or an iterable of [CachedElement][arkindex_worker.cache.CachedElement] when caching is enabled.
         """
         assert element and isinstance(
-            element, (Element, CachedElement)
+            element, Element | CachedElement
         ), "element shouldn't be null and should be an Element or CachedElement"
         query_params = {}
         if folder is not None:
@@ -585,7 +587,7 @@ class ElementMixin(object):
             query_params["recursive"] = recursive
         if transcription_worker_version is not None:
             assert isinstance(
-                transcription_worker_version, (str, bool)
+                transcription_worker_version, str | bool
             ), "transcription_worker_version should be of type str or bool"
             if isinstance(transcription_worker_version, bool):
                 assert (
@@ -594,7 +596,7 @@ class ElementMixin(object):
             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, str | bool
             ), "transcription_worker_run should be of type str or bool"
             if isinstance(transcription_worker_run, bool):
                 assert (
@@ -625,7 +627,7 @@ class ElementMixin(object):
             query_params["with_zone"] = with_zone
         if worker_version is not None:
             assert isinstance(
-                worker_version, (str, bool)
+                worker_version, str | bool
             ), "worker_version should be of type str or bool"
             if isinstance(worker_version, bool):
                 assert (
@@ -634,7 +636,7 @@ class ElementMixin(object):
             query_params["worker_version"] = worker_version
         if worker_run is not None:
             assert isinstance(
-                worker_run, (str, bool)
+                worker_run, str | bool
             ), "worker_run should be of type str or bool"
             if isinstance(worker_run, bool):
                 assert (
@@ -644,11 +646,14 @@ class ElementMixin(object):
 
         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'"
+            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

arkindex_worker/worker/entity.py

@@ -1,10 +1,9 @@
-# -*- coding: utf-8 -*-
 """
 ElementsWorker methods for entities.
 """
 
 from operator import itemgetter
-from typing import Dict, List, Optional, TypedDict, Union
+from typing import TypedDict
 
 from peewee import IntegrityError
 
@@ -12,16 +11,13 @@ from arkindex_worker import logger
 from arkindex_worker.cache import CachedEntity, CachedTranscriptionEntity
 from arkindex_worker.models import Element, Transcription
 
-Entity = TypedDict(
-    "Entity",
-    {
-        "name": str,
-        "type_id": str,
-        "length": int,
-        "offset": int,
-        "confidence": Optional[float],
-    },
-)
+
+class Entity(TypedDict):
+    name: str
+    type_id: str
+    length: int
+    offset: int
+    confidence: float | None
 
 
 class MissingEntityType(Exception):
@@ -31,9 +27,9 @@ class MissingEntityType(Exception):
     """
 
 
-class EntityMixin(object):
+class EntityMixin:
     def check_required_entity_types(
-        self, entity_types: List[str], create_missing: bool = True
+        self, entity_types: list[str], create_missing: bool = True
     ):
         """Checks that every entity type needed is available in the corpus.
         Missing ones may be created automatically if needed.
@@ -71,7 +67,7 @@ class EntityMixin(object):
         self,
         name: str,
         type: str,
-        metas=dict(),
+        metas=None,
         validated=None,
     ):
         """
@@ -87,6 +83,7 @@ class EntityMixin(object):
         assert type and isinstance(
             type, str
         ), "type shouldn't be null and should be of type str"
+        metas = metas or {}
         if metas:
             assert isinstance(metas, dict), "metas should be of type dict"
         if validated is not None:
@@ -140,8 +137,8 @@ class EntityMixin(object):
         entity: str,
         offset: int,
         length: int,
-        confidence: Optional[float] = None,
-    ) -> Optional[Dict[str, Union[str, int]]]:
+        confidence: float | None = None,
+    ) -> dict[str, str | int] | None:
         """
         Create a link between an existing entity and an existing transcription.
         If cache support is enabled, a `CachedTranscriptionEntity` will also be created.
@@ -211,8 +208,8 @@ class EntityMixin(object):
     def create_transcription_entities(
         self,
         transcription: Transcription,
-        entities: List[Entity],
-    ) -> List[Dict[str, str]]:
+        entities: list[Entity],
+    ) -> list[dict[str, str]]:
         """
         Create multiple entities attached to a transcription in a single API request.
 
@@ -250,13 +247,13 @@ class EntityMixin(object):
             ), f"Entity at index {index} in entities: Should be of type dict"
 
             name = entity.get("name")
-            assert name and isinstance(
-                name, str
+            assert (
+                name and isinstance(name, str)
             ), f"Entity at index {index} in entities: name shouldn't be null and should be of type str"
 
             type_id = entity.get("type_id")
-            assert type_id and isinstance(
-                type_id, str
+            assert (
+                type_id and isinstance(type_id, str)
             ), f"Entity at index {index} in entities: type_id shouldn't be null and should be of type str"
 
             offset = entity.get("offset")
@@ -270,8 +267,9 @@ class EntityMixin(object):
             ), f"Entity at index {index} in entities: length shouldn't be null and should be a strictly positive integer"
 
             confidence = entity.get("confidence")
-            assert confidence is None or (
-                isinstance(confidence, float) and 0 <= confidence <= 1
+            assert (
+                confidence is None
+                or (isinstance(confidence, float) and 0 <= confidence <= 1)
             ), f"Entity at index {index} in entities: confidence should be None or a float in [0..1] range"
 
         assert len(entities) == len(
@@ -298,7 +296,7 @@ class EntityMixin(object):
     def list_transcription_entities(
         self,
         transcription: Transcription,
-        worker_version: Optional[Union[str, bool]] = None,
+        worker_version: str | bool | None = None,
     ):
         """
         List existing entities on a transcription
@@ -314,7 +312,7 @@ class EntityMixin(object):
 
         if worker_version is not None:
             assert isinstance(
-                worker_version, (str, bool)
+                worker_version, str | bool
             ), "worker_version should be of type str or bool"
 
             if isinstance(worker_version, bool):
@@ -329,8 +327,8 @@ class EntityMixin(object):
 
     def list_corpus_entities(
         self,
-        name: Optional[str] = None,
-        parent: Optional[Element] = None,
+        name: str | None = None,
+        parent: Element | None = None,
     ):
         """
         List all entities in the worker's corpus