Merge pull request #940 from Labelbox/lgluszek/ontology-update

lgluszek · web-flow · commit 5a796da649e2 · 2023-03-01T19:16:40.000+01:00
New methods to manage feature schemas and ontologies
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
@@ -28,7 +29,7 @@
 from labelbox.schema.labeling_frontend import LabelingFrontend
 from labelbox.schema.model import Model
 from labelbox.schema.model_run import ModelRun
-from labelbox.schema.ontology import Ontology, Tool, Classification
+from labelbox.schema.ontology import Ontology, Tool, Classification, FeatureSchema
 from labelbox.schema.organization import Organization
 from labelbox.schema.user import User
 from labelbox.schema.project import Project
@@ -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,192 @@ 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) -> None:
+        """
+        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) -> None:
+        """
+        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) -> FeatureSchema:
+        """
+        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, feature_schema: Dict) -> FeatureSchema:
+        """
+        Upserts a feature schema
+        Args:
+            feature_schema: Dict representing 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 = feature_schema.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(feature_schema)},
+        )
+
+        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) -> None:
+        """
+        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 get_unused_ontologies(self, after: str = None) -> List[str]:
+        """
+        Returns a list of unused ontology ids
+        Args:
+            after (str): The cursor to use for pagination
+        Returns:
+            A list of unused ontology ids
+        Example:
+            To get the first page of unused ontology ids (100 at a time)
+            >>> client.get_unused_ontologies()
+            To get the next page of unused ontology ids
+            >>> client.get_unused_ontologies("cleabc1my012ioqvu5anyaabc")
+        """
+
+        endpoint = self.rest_endpoint + "/ontologies/unused"
+        response = requests.get(
+            endpoint,
+            headers=self.headers,
+            json={"after": after},
+        )
+
+        if response.status_code == requests.codes.ok:
+            return response.json()
+        else:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to get unused ontologies, message: " +
+                str(response.json()['message']))
+
+    def get_unused_feature_schemas(self, after: str = None) -> List[str]:
+        """
+        Returns a list of unused feature schema ids
+        Args:
+            after (str): The cursor to use for pagination
+        Returns:
+            A list of unused feature schema ids
+        Example:
+            To get the first page of unused feature schema ids (100 at a time)
+            >>> client.get_unused_feature_schemas()
+            To get the next page of unused feature schema ids
+            >>> client.get_unused_feature_schemas("cleabc1my012ioqvu5anyaabc")
+        """
+
+        endpoint = self.rest_endpoint + "/feature-schemas/unused"
+        response = requests.get(
+            endpoint,
+            headers=self.headers,
+            json={"after": after},
+        )
+
+        if response.status_code == requests.codes.ok:
+            return response.json()
+        else:
+            raise labelbox.exceptions.LabelboxError(
+                "Failed to get unused feature schemas, message: " +
+                str(response.json()['message']))
+
     def create_ontology(self, name, normalized, media_type=None) -> Ontology:
         """
         Creates an ontology from normalized data
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -65,6 +65,19 @@ 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:
+        rest_api_endpoint = os.environ.get('LABELBOX_TEST_REST_API_ENDPOINT')
+        if rest_api_endpoint is None:
+            raise Exception(f"Missing LABELBOX_TEST_REST_API_ENDPOINT")
+        return rest_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 +144,11 @@ 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,129 @@
+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"
+
+
+def test_does_not_include_used_feature_schema(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)
+    unused_feature_schemas = client.get_unused_feature_schemas()
+
+    assert feature_schema_id not in unused_feature_schemas
+
+    client.delete_unused_ontology(ontology.uid)
+    client.delete_unused_feature_schema(feature_schema_id)
diff --git a/tests/integration/test_ontology.py b/tests/integration/test_ontology.py
diff --git a/tox.ini b/tox.ini
