Simplify classes ID attribute

Closes #192 (closed)

arkindex_worker/worker/classification.py

@@ -23,10 +23,10 @@ class ClassificationMixin(object):
             "ListCorpusMLClasses",
             id=self.corpus_id,
         )
-        self.classes[self.corpus_id] = {
-            ml_class["name"]: ml_class["id"] for ml_class in corpus_classes
-        }
-        logger.info(f"Loaded {len(self.classes[self.corpus_id])} ML classes")
+        self.classes = {ml_class["name"]: ml_class["id"] for ml_class in corpus_classes}
+        logger.info(
+            f"Loaded {len(self.classes)} ML classes in corpus ({self.corpus_id})"
+        )
 
     def get_ml_class_id(self, ml_class: str) -> str:
         """
@@ -36,17 +36,17 @@ class ClassificationMixin(object):
         :param ml_class: Name of the MLClass.
         :returns: ID of the retrieved or created MLClass.
         """
-        if not self.classes.get(self.corpus_id):
+        if not self.classes:
             self.load_corpus_classes()
 
-        ml_class_id = self.classes[self.corpus_id].get(ml_class)
+        ml_class_id = self.classes.get(ml_class)
         if ml_class_id is None:
             logger.info(f"Creating ML class {ml_class} on corpus {self.corpus_id}")
             try:
                 response = self.request(
                     "CreateMLClass", id=self.corpus_id, body={"name": ml_class}
                 )
-                ml_class_id = self.classes[self.corpus_id][ml_class] = response["id"]
+                ml_class_id = self.classes[ml_class] = response["id"]
                 logger.debug(f"Created ML class {response['id']}")
             except ErrorResponse as e:
                 # Only reload for 400 errors
@@ -59,9 +59,9 @@ class ClassificationMixin(object):
                 )
                 self.load_corpus_classes()
                 assert (
-                    ml_class in self.classes[self.corpus_id]
+                    ml_class in self.classes
                 ), "Missing class {ml_class} even after reloading"
-                ml_class_id = self.classes[self.corpus_id][ml_class]
+                ml_class_id = self.classes[ml_class]
 
         return ml_class_id
 
@@ -73,14 +73,14 @@ class ClassificationMixin(object):
         :return: The MLClass's name
         """
         # Load the corpus' MLclasses if they are not available yet
-        if self.corpus_id not in self.classes:
+        if not self.classes:
             self.load_corpus_classes()
 
         # Filter classes by this ml_class_id
         ml_class_name = next(
             filter(
-                lambda x: self.classes[self.corpus_id][x] == ml_class_id,
-                self.classes[self.corpus_id],
+                lambda x: self.classes[x] == ml_class_id,
+                self.classes,
             ),
             None,
         )