create missing types option

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
 

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
 ):