Add ability to query.count

Author: VaghinakDev

src/superannotate/__init__.py

@@ -2,7 +2,7 @@
 import os
 import sys
 
-__version__ = "4.5.4dev1"
+__version__ = "4.5.5dev2"
 
 
 os.environ.update({"sa_version": __version__})

src/superannotate/lib/app/interface/responses.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from typing import Callable
+from typing import Generic
+from typing import Iterator
+from typing import TypeVar
+from typing import overload
+
+T = TypeVar("T")
+
+
+class BaseResult(Generic[T]):
+    """A generic list-like wrapper for results with lazy loading support.
+
+    This class wraps a list of results while maintaining full backward
+    compatibility with list-like operations (iteration, indexing, len()).
+    Data is fetched lazily on first access.
+    """
+
+    def __init__(self, data_fetcher: Callable[[], list[T]]) -> None:
+        self._data: list[T] | None = None
+        self._data_fetcher = data_fetcher
+
+    def _ensure_data(self) -> list[T]:
+        """Lazily fetch data if not already loaded."""
+        if self._data is None:
+            self._data = self._data_fetcher()
+        return self._data
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self._ensure_data())
+
+    def __len__(self) -> int:
+        return len(self._ensure_data())
+
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> list[T]: ...
+
+    def __getitem__(self, index: int | slice) -> T | list[T]:
+        return self._ensure_data()[index]
+
+    def __repr__(self) -> str:
+        return repr(self._ensure_data())
+
+    def __bool__(self) -> bool:
+        return bool(self._ensure_data())
+
+    def __contains__(self, item: T) -> bool:
+        return item in self._ensure_data()
+
+
+class QueryResult(BaseResult[dict]):
+    """A list-like wrapper for query results that supports .count() method.
+
+    This class wraps a list of query results while maintaining full backward
+    compatibility with list-like operations (iteration, indexing, len()).
+    Data is fetched lazily - only when accessed. Calling .count() does not
+    trigger data fetching.
+    """
+
+    def __init__(
+        self,
+        data_fetcher: Callable[[], list[dict]],
+        count_fetcher: Callable[[], int],
+    ) -> None:
+        super().__init__(data_fetcher)
+        self._count_fetcher = count_fetcher
+
+    def count(self) -> int:
+        """Return the count of items matching the query from the server.
+
+        This method does not trigger data fetching - it makes a separate
+        lightweight API call to get only the count.
+        """
+        return self._count_fetcher()

src/superannotate/lib/app/interface/sdk_interface.py

@@ -11,6 +11,7 @@
 import warnings
 from collections.abc import Callable
 from collections.abc import Iterable
+from functools import partial
 from pathlib import Path
 from typing import Annotated
 from typing import Any
@@ -80,6 +81,8 @@
 from lib.infrastructure.query_builder import QueryBuilderChain
 from lib.infrastructure.query_builder import FieldValidationHandler
 
+from lib.app.interface.responses import QueryResult
+
 logger = logging.getLogger("sa")
 
 NotEmptyStr = Annotated[str, StringConstraints(strict=True, min_length=1)]
@@ -4267,10 +4270,14 @@ def query(
         project: NotEmptyStr | int | tuple[int, int] | tuple[str, str],
         query: NotEmptyStr | None = None,
         subset: NotEmptyStr | None = None,
-    ):
+    ) -> QueryResult:
         """Return items that satisfy the given query.
         Query syntax should be in SuperAnnotate query language(https://doc.superannotate.com/docs/explore-overview).
 
+        The returned object behaves like a list of dicts (supports iteration,
+        indexing, and ``len()``) and additionally exposes a ``.count()`` method
+        that returns the total number of matching items without fetching them.
+
         :param project: Accepts a project as a string ("project" or "project/folder") or as a tuple (project_id, folder_id), where the folder is optional.”
         :type project: Union[str, int, Tuple[int, int], Tuple[str, str]]
 
@@ -4283,13 +4290,39 @@ def query(
 
         :return: queried items' metadata list
         :rtype: list of dicts
+
+        Request Example:
+        ::
+
+            client = SAClient()
+
+            # Iterate over queried items (fetches data)
+            queried_items = client.query(
+                project="Image Project",
+                query="instance(error = true)"
+            )
+            for item in queried_items:
+                print(item["name"])
+
+            # Get only the count without fetching all items
+            total = client.query(
+                project="Image Project",
+                query="instance(error = true)"
+            ).count()
+            print(f"Total matching items: {total}")
         """
-        project, folder = self.controller.get_project_folder(project)
-        items = self.controller.query_entities(project, folder, query, subset)
-        exclude = {
-            "meta",
-        }
-        return BaseSerializer.serialize_iterable(items, exclude=exclude)
+        project_entity, folder = self.controller.get_project_folder(project)
+        fetch_entities = partial(
+            self.controller.query_entities, project_entity, folder, query, subset
+        )
+        return QueryResult(
+            data_fetcher=lambda: BaseSerializer.serialize_iterable(
+                fetch_entities(), exclude={"meta"}
+            ),
+            count_fetcher=partial(
+                self.controller.query_items_count, project_entity.name, query
+            ),
+        )
 
     def get_item_metadata(
         self,

tests/integration/items/test_saqul_query.py

@@ -66,6 +66,80 @@ def test_query_on_100(self):
             == 100
         )
 
+    def test_query_result_list_like_behavior(self):
+        """Test that QueryResult behaves like a list for backward compatibility."""
+        sa.attach_items(self.PROJECT_NAME, os.path.join(DATA_SET_PATH, "100_urls.csv"))
+        result = sa.query(self.PROJECT_NAME, "metadata(status = NotStarted)")
+
+        # Test len()
+        self.assertEqual(len(result), 100)
+
+        # Test indexing
+        first_item = result[0]
+        self.assertIsInstance(first_item, dict)
+        self.assertIn("name", first_item)
+
+        # Test negative indexing
+        last_item = result[-1]
+        self.assertIsInstance(last_item, dict)
+
+        # Test slicing
+        sliced = result[0:5]
+        self.assertEqual(len(sliced), 5)
+
+        # Test iteration
+        count = 0
+        for item in result:
+            self.assertIsInstance(item, dict)
+            count += 1
+        self.assertEqual(count, 100)
+
+        # Test list conversion
+        as_list = list(result)
+        self.assertEqual(len(as_list), 100)
+        self.assertIsInstance(as_list, list)
+
+    def test_query_result_count_method(self):
+        """Test that QueryResult.count() returns the count from server."""
+        sa.attach_items(self.PROJECT_NAME, os.path.join(DATA_SET_PATH, "100_urls.csv"))
+        result = sa.query(self.PROJECT_NAME, "metadata(status = NotStarted)")
+
+        # Test .count() method
+        count = result.count()
+        self.assertEqual(count, 100)
+        self.assertIsInstance(count, int)
+
+        # Verify count matches len
+        self.assertEqual(count, len(result))
+
+    def test_query_result_lazy_loading(self):
+        """Test that QueryResult.count() does not trigger data fetching."""
+        sa.attach_items(self.PROJECT_NAME, os.path.join(DATA_SET_PATH, "100_urls.csv"))
+        result = sa.query(self.PROJECT_NAME, "metadata(status = NotStarted)")
+
+        # Data should not be loaded yet
+        self.assertIsNone(result._data)
+
+        # Calling count() should not load data
+        count = result.count()
+        self.assertEqual(count, 100)
+        self.assertIsNone(result._data)
+
+        # Accessing data should trigger loading
+        first_item = result[0]
+        self.assertIsNotNone(result._data)
+        self.assertIsInstance(first_item, dict)
+
+    def test_query_result_repr(self):
+        """Test that QueryResult repr shows the underlying list."""
+        sa.attach_items(self.PROJECT_NAME, os.path.join(DATA_SET_PATH, "100_urls.csv"))
+        result = sa.query(self.PROJECT_NAME, "metadata(status = NotStarted)")
+
+        # Test __repr__
+        repr_str = repr(result)
+        self.assertIsInstance(repr_str, str)
+        self.assertTrue(repr_str.startswith("["))
+
     def test_validate_saqul_query(self):
         try:
             self.assertRaises(