Added api support for adding a taxonomy and an accompanying test

Author: Claire Pajot

nucleus/__init__.py

@@ -128,6 +128,7 @@
     construct_box_predictions_payload,
     construct_model_creation_payload,
     construct_segmentation_payload,
+    construct_taxonomy_payload,
 )
 from .prediction import (
     BoxPrediction,
@@ -1249,6 +1250,27 @@ def create_image_index(self, dataset_id: str):
             requests_command=requests.post,
         )
 
+    def add_taxonomy(
+        self,
+        dataset_id: str,
+        taxonomy_name: str,
+        taxonomy_type: str,
+        labels: List[str],
+    ):
+        """
+        Creates a new taxonomy.
+        Returns a response with dataset_id, taxonomy_name and type for the new taxonomy.
+        :param dataset_id: id of dataset to add the taxonomy to
+        :param taxonomy_name: name of the taxonomy
+        :param type: type of the taxonomy
+        :param labels: list of possible labels for the taxonomy
+        """
+        return self.make_request(
+            construct_taxonomy_payload(taxonomy_name, taxonomy_type, labels),
+            f"dataset/{dataset_id}/add_taxonomy",
+            requests_command=requests.post,
+        )
+
     def make_request(
         self, payload: dict, route: str, requests_command=requests.post
     ) -> dict:

nucleus/constants.py

@@ -8,7 +8,10 @@
 MASK_TYPE = "mask"
 SEGMENTATION_TYPE = "segmentation"
 CUBOID_TYPE = "cuboid"
+CATEGORY_TYPE = "category"
+MULTICATEGORY_TYPE = "multicategory"
 ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
+TAXONOMY_TYPES = (CATEGORY_TYPE, MULTICATEGORY_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
 EXPORTED_ROWS = "exportedRows"
@@ -58,6 +61,7 @@
 JOB_TYPE_KEY = "job_type"
 JOB_CREATION_TIME_KEY = "job_creation_time"
 LABEL_KEY = "label"
+LABELS_KEY = "labels"
 MASK_URL_KEY = "mask_url"
 MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"
@@ -83,6 +87,7 @@
 STATUS_CODE_KEY = "status_code"
 STATUS_KEY = "status"
 SUCCESS_STATUS_CODES = [200, 201, 202]
+TAXONOMY_NAME_KEY = "taxonomy_name"
 TYPE_KEY = "type"
 UPDATED_ITEMS = "updated_items"
 UPDATE_KEY = "update"

nucleus/dataset.py

@@ -407,6 +407,13 @@ def create_image_index(self):
         response = self._client.create_image_index(self.id)
         return AsyncJob.from_json(response, self._client)
 
+    def add_taxonomy(
+        self, taxonomy_name: str, taxonomy_type: str, labels: List[str]
+    ):
+        return self._client.add_taxonomy(
+            self.id, taxonomy_name, taxonomy_type, labels
+        )
+
     def check_index_status(self, job_id: str):
         return self._client.check_index_status(job_id)
 

nucleus/payload_constructor.py

@@ -25,6 +25,9 @@
     MODEL_ID_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
     SEGMENTATIONS_KEY,
+    TAXONOMY_NAME_KEY,
+    TYPE_KEY,
+    LABELS_KEY,
 )
 
 
@@ -126,3 +129,13 @@ def construct_model_run_creation_payload(
         METADATA_KEY: metadata if metadata else {},
         ANNOTATION_METADATA_SCHEMA_KEY: annotation_metadata_schema,
     }
+
+
+def construct_taxonomy_payload(
+    taxonomy_name: str, taxonomy_type: str, labels: List[str]
+) -> dict:
+    return {
+        TAXONOMY_NAME_KEY: taxonomy_name,
+        TYPE_KEY: taxonomy_type,
+        LABELS_KEY: labels,
+    }

tests/test_taxonomy.py

@@ -0,0 +1,26 @@
+import pytest
+
+from .helpers import (
+    TEST_DATASET_NAME,
+)
+
+
+@pytest.fixture()
+def dataset(CLIENT):
+    ds = CLIENT.create_dataset(TEST_DATASET_NAME)
+    yield ds
+
+    response = CLIENT.delete_dataset(ds.id)
+    assert response == {"message": "Beginning dataset deletion..."}
+
+
+def test_add_taxonomy(dataset):
+    response = dataset.add_taxonomy(
+        "[Pytest] taxonomy",
+        "category",
+        ["[Pytest] taxonomy label 1", "[Pytest] taxonomy label 2"],
+    )
+
+    assert response["dataset_id"] == dataset.id
+    assert response["taxonomy_name"] == "[Pytest] taxonomy"
+    assert response["type"] == "category"