Add a check_required_types method

Closes #73 (closed)

arkindex_worker/worker/element.py

@@ -8,7 +8,34 @@ from arkindex_worker.cache import CachedElement, CachedImage
 from arkindex_worker.models import Element
 
 
+class MissingTypeError(Exception):
+    """
+    A required element type was not found in a corpus.
+    """
+
+
 class ElementMixin(object):
+    def check_required_types(self, corpus_id: str, *type_slugs: str) -> bool:
+        """
+        Check that a corpus has a list of required element types,
+        and raise an exception if any of them are missing.
+        """
+        assert len(type_slugs), "At least one element type slug is required."
+        assert all(
+            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"]}
+        missing_slugs = set(type_slugs) - available_slugs
+
+        if missing_slugs:
+            raise MissingTypeError(
+                f'Element type(s) {", ".join(missing_slugs)} were not found in the {corpus["name"]} corpus ({corpus["id"]}).'
+            )
+
+        return True
+
     def create_sub_element(self, element, type, name, polygon):
         """
         Create a child element on the given element through API