Create missing types option when checking required types existence

Closes #138 (closed)

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("ElementType", name=str, slug=str, is_folder=bool)
 
 
 class MissingTypeError(Exception):
@@ -24,13 +26,34 @@ 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,
+                },
+            )
+            logger.info(f"Created a new element type with slug {element_type.slug}")
+
+    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 +65,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