Add Slice object and function to export data row IDs by slice

Author: attila-papai

labelbox/__init__.py

@@ -27,3 +27,4 @@
 from labelbox.schema.resource_tag import ResourceTag
 from labelbox.schema.project_resource_tag import ProjectResourceTag
 from labelbox.schema.media_type import MediaType
+from labelbox.schema.slice import Slice

labelbox/client.py

@@ -33,6 +33,7 @@
 from labelbox.schema.user import User
 from labelbox.schema.project import Project
 from labelbox.schema.role import Role
+from labelbox.schema.slice import Slice
 
 from labelbox.schema.media_type import MediaType
 
@@ -963,7 +964,7 @@ def assign_global_keys_to_data_rows(
             timeout_seconds=60) -> Dict[str, Union[str, List[Any]]]:
         """
         Assigns global keys to data rows.
-        
+
         Args:
             A list of dicts containing data_row_id and global_key.
         Returns:
@@ -1211,3 +1212,27 @@ def _format_failed_rows(rows: List[str],
                     "Timed out waiting for get_data_rows_for_global_keys job to complete."
                 )
             time.sleep(sleep_time)
+
+    def get_slice(self, slice_id) -> Slice:
+        """
+        Fetches a Slice by ID.
+
+        Args:
+             slice_id (str): The ID of the Slice
+        Returns:
+            Slice
+        """
+        query_str = """
+            query getSavedQueryPyApi($id: ID!) {
+                getSavedQuery(id: $id) {
+                    id
+                    name
+                    description
+                    filter
+                    createdAt
+                    updatedAt
+                }
+            }
+        """
+        res = self.execute(query_str, {'id': slice_id})
+        return Entity.Slice(self, res['getSavedQuery'])

labelbox/schema/slice.py

@@ -0,0 +1,49 @@
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field
+
+
+class Slice(DbObject):
+    name = Field.String("name")
+    description = Field.String("description")
+    created_at = Field.DateTime("created_at")
+    updated_at = Field.DateTime("updated_at")
+    filter = Field.Json("filter")
+
+    def get_data_row_ids(self) -> list[str]:
+        """
+        Fetches all data row ids that match this Slice
+
+        Returns:
+            A list of data row ids
+        """
+        query_str = """
+            query getDataRowIdsBySavedQueryPyApi($id: ID!, $after: String) {
+                getDataRowIdsBySavedQuery(input: {
+                    savedQueryId: $id,
+                    after: $after
+                }) {
+                    totalCount
+                    nodes
+                    pageInfo {
+                        endCursor
+                        hasNextPage
+                    }
+                }
+            }
+        """
+        data_row_ids = []
+        total_count = 0
+        end_cursor = None
+        has_next_page = True
+        while has_next_page:
+            res = self.client.execute(query_str, {
+                'id': self.uid,
+                'after': end_cursor
+            })['getDataRowIdsBySavedQuery']
+            data_row_ids = data_row_ids + res['nodes']
+            total_count = res['totalCount']
+            has_next_page = res['pageInfo']['hasNextPage']
+            end_cursor = res['pageInfo']['endCursor']
+
+        assert total_count == len(data_row_ids)
+        return data_row_ids