arkindex-base-worker 0.3.6rc3__tar.gz → 0.3.6rc4__tar.gz

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.3.6rc3
+Version: 0.3.6rc4
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/arkindex_base_worker.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arkindex-base-worker
-Version: 0.3.6rc3
+Version: 0.3.6rc4
 Summary: Base Worker to easily build Arkindex ML workflows
 Author-email: Teklia <contact@teklia.com>
 Maintainer-email: Teklia <contact@teklia.com>

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/arkindex_worker/worker/__init__.py

@@ -387,7 +387,7 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         """
 
         def format_split(
-            split: tuple[str, Iterator[tuple[str, Element]]]
+            split: tuple[str, Iterator[tuple[str, Element]]],
         ) -> tuple[str, list[Element]]:
             return (split[0], list(map(itemgetter(1), list(split[1]))))
 
@@ -440,6 +440,8 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
         failed = 0
         for i, item in enumerate(datasets, start=1):
             dataset = None
+            dataset_artifact = None
+
             try:
                 if not self.is_read_only:
                     # Just use the result of list_datasets as the dataset
@@ -465,7 +467,7 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
                     self.update_dataset_state(dataset, DatasetState.Building)
                 else:
                     logger.info(f"Downloading data for {dataset} ({i}/{count})")
-                    self.download_dataset_artifact(dataset)
+                    dataset_artifact = self.download_dataset_artifact(dataset)
 
                 # Process the dataset
                 self.process_dataset(dataset)
@@ -496,6 +498,10 @@ class DatasetWorker(BaseWorker, DatasetMixin, TaskMixin):
                     # Try to update the state to Error regardless of the response
                     with contextlib.suppress(Exception):
                         self.update_dataset_state(dataset, DatasetState.Error)
+            finally:
+                # Cleanup the dataset artifact if it was downloaded, no matter what
+                if dataset_artifact:
+                    dataset_artifact.unlink(missing_ok=True)
 
         if failed:
             logger.error(

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/arkindex_worker/worker/entity.py

@@ -331,8 +331,7 @@ class EntityMixin:
         parent: Element | None = None,
     ):
         """
-        List all entities in the worker's corpus
-        This method does not support cache
+        List all entities in the worker's corpus and store them in the ``self.entities`` cache.
         :param name: Filter entities by part of their name (case-insensitive)
         :param parent: Restrict entities to those linked to all transcriptions of an element and all its descendants. Note that links to metadata are ignored.
         """
@@ -346,8 +345,14 @@ class EntityMixin:
             assert isinstance(parent, Element), "parent should be of type Element"
             query_params["parent"] = parent.id
 
-        return self.api_client.paginate(
-            "ListCorpusEntities", id=self.corpus_id, **query_params
+        self.entities = {
+            entity["id"]: entity
+            for entity in self.api_client.paginate(
+                "ListCorpusEntities", id=self.corpus_id, **query_params
+            )
+        }
+        logger.info(
+            f"Loaded {len(self.entities)} entities in corpus ({self.corpus_id})"
         )
 
     def list_corpus_entity_types(

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/pyproject.toml

@@ -4,8 +4,9 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "arkindex-base-worker"
+version = "0.3.6-rc4"
 description = "Base Worker to easily build Arkindex ML workflows"
-dynamic = ["version", "dependencies", "optional-dependencies"]
+dynamic = ["dependencies", "optional-dependencies"]
 authors = [
     { name = "Teklia", email = "contact@teklia.com" },
 ]
@@ -34,8 +35,8 @@ Repository = "https://gitlab.teklia.com/workers/base-worker"
 Authors = "https://teklia.com"
 
 [tool.setuptools.dynamic]
-dependencies = {file = ["requirements.txt"]}
-optional-dependencies = {docs = { file = ["docs-requirements.txt"] }}
+dependencies = { file = ["requirements.txt"] }
+optional-dependencies = { docs = { file = ["docs-requirements.txt"] } }
 
 [tool.ruff]
 exclude = [".git", "__pycache__"]
@@ -52,6 +53,8 @@ select = [
     "I",
     # Implicit Optional
     "RUF013",
+    # Invalid pyproject.toml
+    "RUF200",
     # pyupgrade
     "UP",
     # flake8-bugbear

arkindex-base-worker-0.3.6rc4/setup.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python
+from setuptools import find_packages, setup
+
+setup(packages=find_packages())

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/tests/test_cache.py

@@ -53,6 +53,9 @@ def test_create_tables(tmp_path):
     init_cache_db(db_path)
     create_tables()
 
+    # WARNING: If you are updating this schema following a development you have made
+    # in base-worker, make sure to upgrade the arkindex_worker.cache.SQL_VERSION in
+    # the same merge request as your changes.
     expected_schema = """CREATE TABLE "classifications" ("id" TEXT NOT NULL PRIMARY KEY, "element_id" TEXT NOT NULL, "class_name" TEXT NOT NULL, "confidence" REAL NOT NULL, "state" VARCHAR(10) NOT NULL, "worker_run_id" TEXT, FOREIGN KEY ("element_id") REFERENCES "elements" ("id"))
 CREATE TABLE "dataset_elements" ("id" TEXT NOT NULL PRIMARY KEY, "element_id" TEXT NOT NULL, "dataset_id" TEXT NOT NULL, "set_name" VARCHAR(255) NOT NULL, FOREIGN KEY ("element_id") REFERENCES "elements" ("id"), FOREIGN KEY ("dataset_id") REFERENCES "datasets" ("id"))
 CREATE TABLE "datasets" ("id" TEXT NOT NULL PRIMARY KEY, "name" VARCHAR(255) NOT NULL, "state" VARCHAR(255) NOT NULL DEFAULT 'open', "sets" TEXT NOT NULL)

{arkindex-base-worker-0.3.6rc3 → arkindex-base-worker-0.3.6rc4}/tests/test_elements_worker/test_entities.py

@@ -741,12 +741,13 @@ def test_list_corpus_entities(responses, mock_elements_worker):
         },
     )
 
-    # list is required to actually do the request
-    assert list(mock_elements_worker.list_corpus_entities()) == [
-        {
+    mock_elements_worker.list_corpus_entities()
+
+    assert mock_elements_worker.entities == {
+        "fake_entity_id": {
             "id": "fake_entity_id",
         }
-    ]
+    }
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [

arkindex-base-worker-0.3.6rc3/setup.py

@@ -1,9 +0,0 @@
-#!/usr/bin/env python
-from pathlib import Path
-
-from setuptools import find_packages, setup
-
-setup(
-    version=Path("VERSION").read_text().strip(),
-    packages=find_packages(),
-)