Merge pull request #34 from scaleapi/list-models-endpoint

added ability to list models

Author: srikanth-scale

conftest.py

@@ -17,6 +17,7 @@
 import nucleus
 from nucleus.constants import SUCCESS_STATUS_CODES
 
+from tests.helpers import TEST_DATASET_NAME, TEST_DATASET_ITEMS
 
 assert 'NUCLEUS_PYTEST_API_KEY' in os.environ, \
     "You must set the 'NUCLEUS_PYTEST_API_KEY' environment variable to a valid " \
@@ -55,3 +56,11 @@ def _make_request_patch(
 
     monkeypatch_session.setattr(client, "_make_request", _make_request_patch)
     return client
+
+@pytest.fixture()
+def dataset(CLIENT):
+    ds = CLIENT.create_dataset(TEST_DATASET_NAME)
+    ds.append(TEST_DATASET_ITEMS)
+    yield ds
+
+    CLIENT.delete_dataset(ds.id)

nucleus/__init__.py

@@ -133,13 +133,14 @@ def __init__(self, api_key: str, use_notebook: bool = False):
         if use_notebook:
             self.tqdm_bar = tqdm_notebook.tqdm
 
-    def list_models(self) -> List[str]:
+    def list_models(self) -> List[Model]:
         """
         Lists available models in your repo.
         :return: model_ids
         """
-        # TODO implement API
-        raise NotImplementedError
+        model_objects = self._make_request({}, "models/", requests.get)
+
+        return [Model(model["id"], model["name"], model["ref_id"], model["metadata"], self) for model in model_objects["models"]]
 
     def list_datasets(self) -> Dict[str, Union[str, List[str]]]:
         """

nucleus/model.py

@@ -24,6 +24,15 @@ def __init__(
         self.metadata = metadata
         self._client = client
 
+    def __repr__(self):
+        return f'Model(model_id={self.id}, name={self.name}, reference_id={self.reference_id}, metadata={self.metadata}, client={self._client})'
+
+    def __eq__(self, other):
+        return self.id == other.id
+
+    def __hash__(self):
+        return hash(self.id)
+
     def create_run(
         self,
         name: str,

tests/__init__.py

@@ -1,8 +1,9 @@
 from pathlib import Path
+from nucleus import DatasetItem, BoxPrediction
 
-TEST_MODEL_NAME = '[PyTest] Test Model'
+TEST_MODEL_NAME = '[PyTest] Test Model Name'
 TEST_MODEL_REFERENCE = '[PyTest] Test Model Reference'
-TEST_MODEL_RUN = '[PyTest] Test Model Run'
+TEST_MODEL_RUN = '[PyTest] Test Model Run Reference'
 TEST_DATASET_NAME = '[PyTest] Test Dataset'
 TEST_SLICE_NAME = '[PyTest] Test Slice'
 TEST_IMG_URLS = [
@@ -12,6 +13,18 @@
     's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/06924f46-1708b96f.jpg',
     's3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg',
 ]
+TEST_DATASET_ITEMS = [
+    DatasetItem(TEST_IMG_URLS[0], '1'),
+    DatasetItem(TEST_IMG_URLS[1], '2'),
+    DatasetItem(TEST_IMG_URLS[2], '3'),
+    DatasetItem(TEST_IMG_URLS[3], '4')
+]
+TEST_PREDS = [
+    BoxPrediction('[Pytest Box Prediction 1]', 0, 0, 100, 100, '1'),
+    BoxPrediction('[Pytest Box Prediction 2]', 0, 0, 100, 100, '2'),
+    BoxPrediction('[Pytest Box Prediction 3]', 0, 0, 100, 100, '3'),
+    BoxPrediction('[Pytest Box Prediction 4]', 0, 0, 100, 100, '4')
+]
 
 def reference_id_from_url(url):
     return Path(url).name

tests/helpers.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+import pytest
+from nucleus import (
+    Dataset,
+    DatasetItem,
+    UploadResponse,
+    Model,
+    ModelRun,
+    BoxPrediction
+)
+from nucleus.constants import (
+    NEW_ITEMS,
+    UPDATED_ITEMS,
+    IGNORED_ITEMS,
+    ERROR_ITEMS,
+    ERROR_PAYLOAD,
+    DATASET_ID_KEY,
+)
+from .helpers import (
+    TEST_MODEL_NAME,
+    TEST_MODEL_REFERENCE,
+    TEST_MODEL_RUN,
+    TEST_PREDS
+)
+
+def test_model_creation_and_listing(CLIENT, dataset):
+    models_before = CLIENT.list_models()
+
+    # Creation
+    model = CLIENT.add_model(TEST_MODEL_NAME, TEST_MODEL_REFERENCE)
+    m_run =  model.create_run(TEST_MODEL_RUN, dataset, TEST_PREDS)
+    m_run.commit()
+
+    assert isinstance(model, Model)
+    assert isinstance(m_run, ModelRun)
+
+    # List the models
+    ms = CLIENT.list_models()
+
+    assert model in ms
+    assert list(set(ms) - set(models_before))[0] == model
+
+    # Delete the model
+    CLIENT.delete_model(model.id)
+    ms = CLIENT.list_models()
+
+    assert model not in ms
+    assert ms == models_before
+