Dataset, DatasetItem, and Client

Author: ardila

nucleus/__init__.py

@@ -141,6 +141,7 @@
 )
 
 
+
 class NucleusClient:
     """
     Nucleus client.
@@ -149,9 +150,21 @@ class NucleusClient:
     def __init__(self, api_key: str, use_notebook: bool = False):
         self.api_key = api_key
         self.tqdm_bar = tqdm.tqdm
+        self._use_notebook = use_notebook
         if use_notebook:
             self.tqdm_bar = tqdm_notebook.tqdm
 
+    def __repr__(self):
+        return f"NucleusClient(api_key='{self.api_key}', use_notebook={self._use_notebook})"
+
+    def __eq__(self, other):
+        if self.api_key == other.api_key:
+            if self._use_notebook == other._use_notebook:
+                return True
+        return False
+
+    
+
     def list_models(self) -> List[Model]:
         """
         Lists available models in your repo.

nucleus/dataset.py

@@ -32,6 +32,15 @@ def __init__(self, dataset_id: str, client):
         self.id = dataset_id
         self._client = client
 
+    def __repr__(self):
+        return f"Dataset(dataset_id='{self.id}', client={self._client})"
+    
+    def __eq__(self, other):
+        if self.id == other.id:
+            if self._client == other._client:
+                return True
+        return False
+
     @property
     def name(self) -> str:
         return self.info().get(DATASET_NAME_KEY, "")

tests/helpers.py

@@ -185,3 +185,7 @@ def assert_polygon_prediction_matches_dict(
         prediction_instance, prediction_dict
     )
     assert prediction_instance.confidence == prediction_dict["confidence"]
+
+
+def test_repr(test_object: any):
+    assert eval(str(test_object)) == test_object

tests/test_annotation.py

@@ -21,11 +21,10 @@
 )
 from nucleus.constants import ERROR_PAYLOAD
 
-
+# Have to define here in order to have access to all relevant objects
 def test_repr(test_object: any):
     assert eval(str(test_object)) == test_object
 
-
 def test_reprs():
     [
         test_repr(SegmentationAnnotation.from_json(_))

tests/test_dataset.py

@@ -8,7 +8,7 @@
     reference_id_from_url,
 )
 
-from nucleus import Dataset, DatasetItem, UploadResponse
+from nucleus import Dataset, DatasetItem, UploadResponse, NucleusClient
 from nucleus.constants import (
     NEW_ITEMS,
     UPDATED_ITEMS,
@@ -19,6 +19,28 @@
 )
 
 
+# Have to define here in order to have access to all relevant objects
+def test_repr(test_object: any):
+    assert eval(str(test_object)) == test_object
+
+def test_reprs():
+    test_repr(DatasetItem(
+                image_location="test_url",
+                reference_id="test_reference_id",
+                metadata={
+                    "made_with_pytest": True,
+                    "example_int": 0,
+                    "example_str": "hello",
+                    "example_float": 0.5,
+                    "example_dict": {
+                        "nested": True,
+                    },
+                    "example_list": ["hello", 1, False],
+                },
+            ))
+    test_repr(Dataset("test_dataset", NucleusClient(api_key="fake_key")))
+
+
 @pytest.fixture()
 def dataset(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)