[IO-1626][external] Unexpected data from API causes crash (#656)

* Initial fix

* Tests for the client fix

* Boyscouting/type clearup

* Changed JSONFreeForm fixes to UnknownType as it is more appropriate.

* Response to PR comments

---------

Co-authored-by: Owen <[email protected]>

Author: owencjones

darwin/client.py

@@ -36,6 +36,7 @@
     is_project_dir,
     urljoin,
 )
+from darwin.utils.get_item_count import get_item_count
 
 
 class Client:
@@ -134,7 +135,7 @@ def list_remote_datasets(self, team_slug: Optional[str] = None) -> Iterator[Remo
                     slug=dataset["slug"],
                     team=team_slug or self.default_team,
                     dataset_id=dataset["id"],
-                    item_count=self._get_item_count(dataset),
+                    item_count=get_item_count(dataset),
                     progress=dataset["progress"],
                     client=self,
                 )
@@ -144,7 +145,7 @@ def list_remote_datasets(self, team_slug: Optional[str] = None) -> Iterator[Remo
                     slug=dataset["slug"],
                     team=team_slug or self.default_team,
                     dataset_id=dataset["id"],
-                    item_count=self._get_item_count(dataset),
+                    item_count=get_item_count(dataset),
                     progress=dataset["progress"],
                     client=self,
                 )
@@ -199,7 +200,7 @@ def get_remote_dataset(self, dataset_identifier: Union[str, DatasetIdentifier])
                     slug=dataset["slug"],
                     team=parsed_dataset_identifier.team_slug,
                     dataset_id=dataset["id"],
-                    item_count=self._get_item_count(dataset),
+                    item_count=get_item_count(dataset),
                     progress=0,
                     client=self,
                 )
@@ -209,7 +210,7 @@ def get_remote_dataset(self, dataset_identifier: Union[str, DatasetIdentifier])
                     slug=dataset["slug"],
                     team=parsed_dataset_identifier.team_slug,
                     dataset_id=dataset["id"],
-                    item_count=self._get_item_count(dataset),
+                    item_count=get_item_count(dataset),
                     progress=0,
                     client=self,
                 )
@@ -245,7 +246,7 @@ def create_dataset(self, name: str, team_slug: Optional[str] = None) -> RemoteDa
                 team=team_slug or self.default_team,
                 slug=dataset["slug"],
                 dataset_id=dataset["id"],
-                item_count=self._get_item_count(dataset),
+                item_count=get_item_count(dataset),
                 progress=0,
                 client=self,
             )
@@ -255,7 +256,7 @@ def create_dataset(self, name: str, team_slug: Optional[str] = None) -> RemoteDa
                 team=team_slug or self.default_team,
                 slug=dataset["slug"],
                 dataset_id=dataset["id"],
-                item_count=self._get_item_count(dataset),
+                item_count=get_item_count(dataset),
                 progress=0,
                 client=self,
             )

darwin/utils/get_item_count.py

@@ -0,0 +1,25 @@
+from typing import Dict, Optional
+
+
+def get_item_count(dataset_dict: Dict) -> int:
+    """
+    Returns the number of items in the dataset.
+
+    Parameters
+    ----------
+    dataset_dict: Dict
+        The dataset dictionary.
+
+    Returns
+    -------
+    int
+        The number of items in the dataset.
+    """
+    num_items: Optional[int] = dataset_dict.get("num_items")
+    num_videos: Optional[int] = dataset_dict.get("num_videos")
+    num_images: Optional[int] = dataset_dict.get("num_images")
+
+    if num_items is not None:
+        return num_items
+
+    return (num_images or 0) + (num_videos or 0)

tests/darwin/utils/get_items_count_test.py

@@ -0,0 +1,41 @@
+# Tests for _get_item_count
+from darwin.utils.get_item_count import get_item_count
+
+
+def test__get_item_count_defaults_to_num_items_if_present() -> None:
+    dataset_return = {
+        "num_images": 2,  # Should be ignored
+        "num_videos": 3,  # Should be ignored
+        "num_items": 5,  # Should get this one
+    }
+
+    assert get_item_count(dataset_return) == 5
+
+
+def test__get_item_count_returns_sum_of_others_if_num_items_not_present() -> None:
+    dataset_return = {
+        "num_images": 7,  # Should be summed
+        "num_videos": 3,  # Should be summed
+    }
+
+    assert get_item_count(dataset_return) == 10
+
+
+def test__get_item_count_should_tolerate_missing_members() -> None:
+    assert (
+        get_item_count(
+            {
+                "num_videos": 3,  # Should be ignored
+            }
+        )
+        == 3
+    )
+
+    assert (
+        get_item_count(
+            {
+                "num_images": 2,
+            }
+        )
+        == 2
+    )