From f4b50a4bae0cb1aae101053c14058ff8756ada38 Mon Sep 17 00:00:00 2001
From: Yoann Schneider <yschneider@teklia.com>
Date: Tue, 30 Aug 2022 17:35:22 +0200
Subject: [PATCH] create missing types option
---
arkindex_worker/worker/element.py | 46 ++++++++++++++----
tests/test_elements_worker/test_elements.py | 53 +++++++++++++++++++++
2 files changed, 91 insertions(+), 8 deletions(-)
diff --git a/arkindex_worker/worker/element.py b/arkindex_worker/worker/element.py
index a22dbd0e..0f3648c0 100644
--- a/arkindex_worker/worker/element.py
+++ b/arkindex_worker/worker/element.py
@@ -4,13 +4,15 @@ ElementsWorker methods for elements and element types.
"""
import uuid
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Dict, Iterable, List, NamedTuple, Optional, Union
from peewee import IntegrityError
from arkindex_worker import logger
from arkindex_worker.cache import CachedElement, CachedImage
-from arkindex_worker.models import Element
+from arkindex_worker.models import Corpus, Element
+
+ElementType = NamedTuple("ElemenType", name=str, slug=str, is_folder=bool)
class MissingTypeError(Exception):
@@ -24,13 +26,33 @@ class ElementMixin(object):
Mixin for the :class:`ElementsWorker` to provide ``Element`` helpers.
"""
- def check_required_types(self, corpus_id: str, *type_slugs: str) -> bool:
+ def create_required_types(self, corpus: Corpus, element_types: List[ElementType]):
+ """Creates given element types in the corpus.
+
+ :param Corpus corpus: The corpus to create types on.
+ :param List[ElementType] element_types: The missing element types to create.
+ """
+ for element_type in element_types:
+ self.request(
+ "CreateElementType",
+ body={
+ "slug": element_type.slug,
+ "display_name": element_type.name,
+ "folder": element_type.is_folder,
+ "corpus": corpus.id,
+ },
+ )
+
+ def check_required_types(
+ self, corpus_id: str, *type_slugs: str, create_missing: bool = False
+ ) -> bool:
"""
Check that a corpus has a list of required element types,
and raise an exception if any of them are missing.
:param str corpus_id: ID of the corpus to check types on.
:param str \\*type_slugs: Type slugs to look for.
+ :param bool create_missing: Whether missing types should be created.
:returns bool: True if all of the specified type slugs have been found.
:raises MissingTypeError: If any of the specified type slugs were not found.
"""
@@ -42,14 +64,22 @@ class ElementMixin(object):
isinstance(slug, str) for slug in type_slugs
), "Element type slugs must be strings."
- corpus = self.request("RetrieveCorpus", id=corpus_id)
- available_slugs = {element_type["slug"] for element_type in corpus["types"]}
+ corpus = Corpus(self.request("RetrieveCorpus", id=corpus_id))
+ available_slugs = {element_type.slug for element_type in corpus.types}
missing_slugs = set(type_slugs) - available_slugs
if missing_slugs:
- raise MissingTypeError(
- f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
- )
+ if create_missing:
+ self.create_required_types(
+ corpus,
+ element_types=[
+ ElementType(slug, slug, False) for slug in missing_slugs
+ ],
+ )
+ else:
+ raise MissingTypeError(
+ f'Element type(s) {", ".join(sorted(missing_slugs))} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
+ )
return True
diff --git a/tests/test_elements_worker/test_elements.py b/tests/test_elements_worker/test_elements.py
index 7ff16855..d3cb20cf 100644
--- a/tests/test_elements_worker/test_elements.py
+++ b/tests/test_elements_worker/test_elements.py
@@ -5,6 +5,7 @@ from uuid import UUID
import pytest
from apistar.exceptions import ErrorResponse
+from responses import matchers
from arkindex_worker.cache import (
SQL_VERSION,
@@ -62,6 +63,58 @@ def test_check_required_types(monkeypatch, tmp_path, mock_elements_worker, respo
)
+def test_create_missing_types(monkeypatch, tmp_path, mock_elements_worker, responses):
+ elements_path = tmp_path / "elements.json"
+ elements_path.write_text("[]")
+ monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+
+ corpus_id = "12341234-1234-1234-1234-123412341234"
+
+ responses.add(
+ responses.GET,
+ f"http://testserver/api/v1/corpus/{corpus_id}/",
+ json={
+ "id": corpus_id,
+ "name": "Some Corpus",
+ "types": [{"slug": "folder"}, {"slug": "page"}],
+ },
+ )
+ responses.add(
+ responses.POST,
+ "http://testserver/api/v1/elements/type/",
+ match=[
+ matchers.json_params_matcher(
+ {
+ "slug": "text_line",
+ "display_name": "text_line",
+ "folder": False,
+ "corpus": corpus_id,
+ }
+ )
+ ],
+ )
+ responses.add(
+ responses.POST,
+ "http://testserver/api/v1/elements/type/",
+ match=[
+ matchers.json_params_matcher(
+ {
+ "slug": "act",
+ "display_name": "act",
+ "folder": False,
+ "corpus": corpus_id,
+ }
+ )
+ ],
+ )
+ worker = ElementsWorker()
+ worker.configure()
+
+ assert worker.check_required_types(
+ corpus_id, "page", "text_line", "act", create_missing=True
+ )
+
+
def test_list_elements_elements_list_arg_wrong_type(
monkeypatch, tmp_path, mock_elements_worker
):
--
GitLab