add methods to manage feature schemas

lgluszek · lgluszek · commit f0f8467f84da · 2023-02-27T13:41:51.000+01:00
diff --git a/Makefile b/Makefile
@@ -44,4 +44,5 @@ test-custom: build
 		-e DA_GCP_LABELBOX_API_KEY=${DA_GCP_LABELBOX_API_KEY} \
 		-e LABELBOX_TEST_API_KEY_CUSTOM=${LABELBOX_TEST_API_KEY_CUSTOM} \
 		-e LABELBOX_TEST_GRAPHQL_API_ENDPOINT=${LABELBOX_TEST_GRAPHQL_API_ENDPOINT} \
+		-e LABELBOX_TEST_REST_API_ENDPOINT=${LABELBOX_TEST_REST_API_ENDPOINT} \
 		local/labelbox-python:test pytest $(PATH_TO_TEST)
diff --git a/labelbox/client.py b/labelbox/client.py
@@ -8,6 +8,7 @@
 import mimetypes
 import os
 import time
+import urllib.parse
 
 from google.api_core import retry
 import requests
@@ -55,7 +56,8 @@ def __init__(self,
                  api_key=None,
                  endpoint='https://api.labelbox.com/graphql',
                  enable_experimental=False,
-                 app_url="https://app.labelbox.com"):
+                 app_url="https://app.labelbox.com",
+                 rest_endpoint="https://api.labelbox.com/api/v1"):
         """ Creates and initializes a Labelbox Client.
 
         Logging is defaulted to level WARNING. To receive more verbose
@@ -88,6 +90,7 @@ def __init__(self,
         logger.info("Initializing Labelbox client at '%s'", endpoint)
         self.app_url = app_url
         self.endpoint = endpoint
+        self.rest_endpoint = rest_endpoint
         self.headers = {
             'Accept': 'application/json',
             'Content-Type': 'application/json',
@@ -899,6 +902,129 @@ def create_ontology_from_feature_schemas(self,
         normalized = {'tools': tools, 'classifications': classifications}
         return self.create_ontology(name, normalized, media_type)
 
+    def delete_unused_feature_schema(self, feature_schema_id: str):
+        """
+               Deletes a feature schema if it is not used by any ontologies or annotations
+               Args:
+                   feature_schema_id (str): The id of the feature schema to delete
+               Example:
+                   >>> client.delete_unused_feature_schema("cleabc1my012ioqvu5anyaabc")
+               """
+
+        endpoint = self.rest_endpoint + "/feature-schemas/" + urllib.parse.quote(feature_schema_id)
+        response = requests.delete(
+            endpoint,
+            headers=self.headers,
+        )
+
+        if response.status_code != requests.codes.no_content:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to delete the feature schema, message: " +
+                str(response.json()['message']))
+
+    def delete_unused_ontology(self, ontology_id: str):
+        """
+        Deletes an ontology if it is not used by any annotations
+        Args:
+            ontology_id (str): The id of the ontology to delete
+        Example:
+            >>> client.delete_unused_ontology("cleabc1my012ioqvu5anyaabc")
+        """
+
+        endpoint = self.rest_endpoint + "/ontologies/" + urllib.parse.quote(ontology_id)
+        response = requests.delete(
+            endpoint,
+            headers=self.headers,
+        )
+
+        if response.status_code != requests.codes.no_content:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to delete the ontology, message: " +
+                str(response.json()['message']))
+
+    def update_feature_schema_title(self, feature_schema_id: str, title: str):
+        """
+        Updates a title of a feature schema
+        Args:
+            feature_schema_id (str): The id of the feature schema to update
+            title (str): The new title of the feature schema
+        Returns:
+            The updated feature schema
+        Example:
+            >>> client.update_feature_schema_title("cleabc1my012ioqvu5anyaabc", "New Title")
+        """
+
+        endpoint = self.rest_endpoint + "/feature-schemas/" + urllib.parse.quote(feature_schema_id) + '/definition'
+        response = requests.patch(
+            endpoint,
+            headers=self.headers,
+            json={"title": title},
+        )
+
+        if response.status_code == requests.codes.ok:
+            return self.get_feature_schema(feature_schema_id)
+        else:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to update the feature schema, message: " +
+                str(response.json()['message']))
+
+    def upsert_feature_schema(self, normalized: Dict):
+        """
+        Upserts a feature schema
+        Args:
+            normalized: The feature schema to upsert
+        Returns:
+            The upserted feature schema
+        Example:
+            Insert a new feature schema
+            >>> tool = Tool(name="tool", tool=Tool.Type.BOUNDING_BOX, color="#FF0000")
+            >>> client.upsert_feature_schema(tool.asdict())
+            Update an existing feature schema
+            >>> tool = Tool(feature_schema_id="cleabc1my012ioqvu5anyaabc", name="tool", tool=Tool.Type.BOUNDING_BOX, color="#FF0000")
+            >>> client.upsert_feature_schema(tool.asdict())
+        """
+
+        feature_schema_id = normalized.get(
+            "featureSchemaId") or "new_feature_schema_id"
+        endpoint = self.rest_endpoint + "/feature-schemas/" + urllib.parse.quote(
+            feature_schema_id)
+        response = requests.put(
+            endpoint,
+            headers=self.headers,
+            json={"normalized": json.dumps(normalized)},
+        )
+
+        if response.status_code == requests.codes.ok:
+            return self.get_feature_schema(response.json()['schemaId'])
+        else:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to upsert the feature schema, message: " +
+                str(response.json()['message']))
+
+    def insert_feature_schema_into_ontology(self, feature_schema_id: str, ontology_id: str, position: int):
+        """
+        Inserts a feature schema into an ontology. If the feature schema is already in the ontology,
+        it will be moved to the new position.
+        Args:
+            feature_schema_id (str): The feature schema id to upsert
+            ontology_id (str): The id of the ontology to insert the feature schema into
+            position (int): The position number of the feature schema in the ontology
+        Example:
+            >>> client.insert_feature_schema_into_ontology("cleabc1my012ioqvu5anyaabc", "clefdvwl7abcgefgu3lyvcde", 2)
+        """
+
+        endpoint = self.rest_endpoint + '/ontologies/' + urllib.parse.quote(
+            ontology_id) + "/feature-schemas/" + urllib.parse.quote(feature_schema_id)
+        response = requests.post(
+            endpoint,
+            headers=self.headers,
+            json={"position": position},
+        )
+        if response.status_code != requests.codes.created:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to insert the feature schema into the ontology, message: " +
+                str(response.json()['message']))
+
     def create_ontology(self, name, normalized, media_type=None) -> Ontology:
         """
         Creates an ontology from normalized data
diff --git a/pytest.ini b/pytest.ini
@@ -1,4 +1,4 @@
 [pytest]
-addopts = -s -vv --reruns 5 --reruns-delay 10 --durations=20
+addopts = -s -vv --reruns 1 --reruns-delay 10 --durations=20
 markers =
     slow: marks tests as slow (deselect with '-m "not slow"')
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -65,6 +65,20 @@ def graphql_url(environ: str) -> str:
     return 'http://host.docker.internal:8080/graphql'
 
 
+def rest_url(environ: str) -> str:
+    if environ == Environ.PROD:
+        return 'https://api.labelbox.com/api/v1'
+    elif environ == Environ.STAGING:
+        return 'https://api.lb-stage.xyz/api/v1'
+    elif environ == Environ.CUSTOM:
+        graphql_api_endpoint = os.environ.get(
+            'LABELBOX_TEST_REST_API_ENDPOINT')
+        if graphql_api_endpoint is None:
+            raise Exception(f"Missing LABELBOX_TEST_REST_API_ENDPOINT")
+        return graphql_api_endpoint
+    return 'http://host.docker.internal:8080/api/v1'
+
+
 def testing_api_key(environ: str) -> str:
     if environ == Environ.PROD:
         return os.environ["LABELBOX_TEST_API_KEY_PROD"]
@@ -131,7 +145,8 @@ class IntegrationClient(Client):
     def __init__(self, environ: str) -> None:
         api_url = graphql_url(environ)
         api_key = testing_api_key(environ)
-        super().__init__(api_key, api_url, enable_experimental=True)
+        rest_endpoint = rest_url(environ)
+        super().__init__(api_key, api_url, enable_experimental=True, rest_endpoint=rest_endpoint)
         self.queries = []
 
     def execute(self, query=None, params=None, check_naming=True, **kwargs):
diff --git a/tests/integration/test_feature_schema.py b/tests/integration/test_feature_schema.py
@@ -0,0 +1,101 @@
+import pytest
+
+from labelbox import Tool, MediaType
+
+point = Tool(
+    tool=Tool.Type.POINT,
+    name="name",
+    color="#ff0000",
+)
+
+
+def test_deletes_a_feature_schema(client):
+    tool = client.upsert_feature_schema(point.asdict())
+
+    assert client.delete_unused_feature_schema(tool.normalized['featureSchemaId']) is None
+
+
+def test_cant_delete_already_deleted_feature_schema(client):
+    tool = client.upsert_feature_schema(point.asdict())
+    feature_schema_id = tool.normalized['featureSchemaId']
+
+    client.delete_unused_feature_schema(feature_schema_id) is None
+
+    with pytest.raises(Exception,
+                       match="Failed to delete the feature schema, message: Feature schema is already deleted"):
+        client.delete_unused_feature_schema(feature_schema_id)
+
+
+def test_cant_delete_feature_schema_with_ontology(client):
+    tool = client.upsert_feature_schema(point.asdict())
+    feature_schema_id = tool.normalized['featureSchemaId']
+    ontology = client.create_ontology_from_feature_schemas(
+        name='ontology name',
+        feature_schema_ids=[feature_schema_id],
+        media_type=MediaType.Image)
+
+    with pytest.raises(Exception,
+                       match="Failed to delete the feature schema, message: Feature schema cannot be deleted because it is used in ontologies"):
+        client.delete_unused_feature_schema(feature_schema_id)
+
+    client.delete_unused_ontology(ontology.uid)
+    client.delete_unused_feature_schema(feature_schema_id)
+
+
+def test_throws_an_error_if_feature_schema_to_delete_doesnt_exist(client):
+    with pytest.raises(Exception,
+                       match="Failed to delete the feature schema, message: Cannot find root schema node with feature schema id doesntexist"):
+        client.delete_unused_feature_schema("doesntexist")
+
+
+def test_updates_a_feature_schema_title(client):
+    tool = client.upsert_feature_schema(point.asdict())
+    feature_schema_id = tool.normalized['featureSchemaId']
+    new_title = "new title"
+    updated_feature_schema = client.update_feature_schema_title(feature_schema_id, new_title)
+
+    assert updated_feature_schema.normalized['name'] == new_title
+
+    client.delete_unused_feature_schema(feature_schema_id)
+
+
+def test_throws_an_error_when_updating_a_feature_schema_with_empty_title(client):
+    tool = client.upsert_feature_schema(point.asdict())
+    feature_schema_id = tool.normalized['featureSchemaId']
+
+    with pytest.raises(Exception):
+        client.update_feature_schema_title(feature_schema_id, "")
+
+    client.delete_unused_feature_schema(feature_schema_id)
+
+
+def test_throws_an_error_when_updating_not_existing_feature_schema(client):
+    with pytest.raises(Exception):
+        client.update_feature_schema_title("doesntexist", "new title")
+
+
+def test_creates_a_new_feature_schema(client):
+    created_feature_schema = client.upsert_feature_schema(point.asdict())
+
+    assert created_feature_schema.uid is not None
+
+    client.delete_unused_feature_schema(created_feature_schema.normalized['featureSchemaId'])
+
+
+def test_updates_a_feature_schema(client):
+    tool = Tool(
+        tool=Tool.Type.POINT,
+        name="name",
+        color="#ff0000",
+    )
+    created_feature_schema = client.upsert_feature_schema(tool.asdict())
+    tool_to_update = Tool(
+        tool=Tool.Type.POINT,
+        name="new name",
+        color="#ff0000",
+        feature_schema_id=created_feature_schema.normalized['featureSchemaId'],
+    )
+    updated_feature_schema = client.upsert_feature_schema(tool_to_update.asdict())
+
+    assert updated_feature_schema.normalized[
+               'name'] == "new name"
diff --git a/tests/integration/test_ontology.py b/tests/integration/test_ontology.py
diff --git a/tox.ini b/tox.ini
