[PY-601][PY-644] Advanced Filters support for ItemQuery (#768)

darwin/future/core/items/get.py

@@ -6,7 +6,7 @@
 from pydantic import ValidationError
 
 from darwin.future.core.client import ClientCore
-from darwin.future.core.types.common import QueryString
+from darwin.future.core.types.common import JSONDict, QueryString
 from darwin.future.data_objects.item import Folder, ItemCore
 
 
@@ -191,3 +191,42 @@ def list_folders(
         except ValidationError as e:
             exceptions.append(e)
     return folders, exceptions
+
+
+def list_items_unstable(
+    api_client: ClientCore,
+    team_slug: str,
+    params: JSONDict,
+) -> Tuple[List[ItemCore], List[ValidationError]]:
+    """
+    Returns a list of items for the dataset from the advanced filters 'unstable' endpoint
+
+    Parameters
+    ----------
+    client: Client
+        The client to use for the request
+    team_slug: str
+        The slug of the team to get items for
+    params: JSONType
+        Must include at least dataset_ids
+
+    Returns
+    -------
+    List[Item]
+        A list of items
+    List[ValidationError]
+        A list of ValidationError on failed objects
+    """
+    if "dataset_ids" not in params:
+        raise ValueError("dataset_ids must be provided")
+    response = api_client.post(f"/unstable/teams/{team_slug}/items/list", params)
+    assert isinstance(response, dict)
+    items: List[ItemCore] = []
+    exceptions: List[ValidationError] = []
+    for item in response["items"]:
+        assert isinstance(item, dict)
+        try:
+            items.append(ItemCore.model_validate(item))
+        except ValidationError as e:
+            exceptions.append(e)
+    return items, exceptions

darwin/future/core/types/query.py

@@ -5,9 +5,11 @@
 from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar
 
 from pydantic import field_validator
+from typing_extensions import Self
 
 from darwin.future.core.client import ClientCore
 from darwin.future.core.types.common import Stringable
+from darwin.future.data_objects.advanced_filters import GroupFilter, SubjectFilter
 from darwin.future.data_objects.page import Page
 from darwin.future.exceptions import (
     InvalidQueryFilter,
@@ -136,9 +138,9 @@ class Query(Generic[T], ABC):
         _changed_since_last (bool): A boolean indicating whether the query has changed since the last execution.
 
     Methods:
-        filter(name: str, param: str, modifier: Optional[Modifier] = None) -> Query[T]:
+        filter(name: str, param: str, modifier: Optional[Modifier] = None) -> Self:
             Adds a filter to the query object and returns a new query object.
-        where(name: str, param: str, modifier: Optional[Modifier] = None) -> Query[T]:
+        where(name: str, param: str, modifier: Optional[Modifier] = None) -> Self:
             Applies a filter on the query object and returns a new query object.
         first() -> Optional[T]:
             Returns the first result of the query. Raises an exception if no results are found.
@@ -172,28 +174,28 @@ def __init__(
         self.results: dict[int, T] = {}
         self._changed_since_last: bool = False
 
-    def filter(self, filter: QueryFilter) -> Query[T]:
+    def filter(self, filter: QueryFilter) -> Self:
         return self + filter
 
-    def __add__(self, filter: QueryFilter) -> Query[T]:
+    def __add__(self, filter: QueryFilter) -> Self:
         self._changed_since_last = True
         self.filters.append(filter)
         return self
 
-    def __sub__(self, filter: QueryFilter) -> Query[T]:
+    def __sub__(self, filter: QueryFilter) -> Self:
         self._changed_since_last = True
         return self.__class__(
             self.client,
             filters=[f for f in self.filters if f != filter],
             meta_params=self.meta_params,
         )
 
-    def __iadd__(self, filter: QueryFilter) -> Query[T]:
+    def __iadd__(self, filter: QueryFilter) -> Self:
         self.filters.append(filter)
         self._changed_since_last = True
         return self
 
-    def __isub__(self, filter: QueryFilter) -> Query[T]:
+    def __isub__(self, filter: QueryFilter) -> Self:
         self.filters = [f for f in self.filters if f != filter]
         self._changed_since_last = True
         return self
@@ -203,7 +205,7 @@ def __len__(self) -> int:
             self.results = {**self.results, **self._collect()}
         return len(self.results)
 
-    def __iter__(self) -> Query[T]:
+    def __iter__(self) -> Self:
         self.n = 0
         return self
 
@@ -227,7 +229,7 @@ def __setitem__(self, index: int, value: T) -> None:
             self.results = {**self.results, **self._collect()}
         self.results[index] = value
 
-    def where(self, *args: object, **kwargs: str) -> Query[T]:
+    def where(self, *args: object, **kwargs: str) -> Self:
         filters = QueryFilter._from_args(*args, **kwargs)
         for item in filters:
             self += item
@@ -280,6 +282,7 @@ def __init__(
         page: Page | None = None,
     ):
         super().__init__(client, filters, meta_params)
+        self._advanced_filters: GroupFilter | SubjectFilter | None = None
         self.page = page or Page()
         self.completed = False
 

darwin/future/data_objects/advanced_filters.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Generic, List, Literal, Optional, TypeVar, Union
 
-from pydantic import BaseModel, field_validator, model_validator
+from pydantic import BaseModel, SerializeAsAny, field_validator, model_validator
 
 T = TypeVar("T")
 
@@ -88,12 +88,16 @@ class NotContains(BaseMatcher):
 
 class SubjectFilter(BaseModel):
     subject: str
-    matcher: BaseMatcher
+    matcher: SerializeAsAny[BaseMatcher]
 
-    def __and__(self, other: SubjectFilter) -> GroupFilter:
+    def __and__(self, other: SubjectFilter | GroupFilter) -> GroupFilter:
+        if isinstance(other, GroupFilter):
+            return GroupFilter(conjunction="and", filters=[self, other])
         return GroupFilter(conjunction="and", filters=[self, other])
 
-    def __or__(self, other: SubjectFilter) -> GroupFilter:
+    def __or__(self, other: SubjectFilter | GroupFilter) -> GroupFilter:
+        if isinstance(other, GroupFilter):
+            return GroupFilter(conjunction="or", filters=[self, other])
         return GroupFilter(conjunction="or", filters=[self, other])
 
 
@@ -301,6 +305,52 @@ def none_of(cls, values: list[ProcessingStatusType]) -> ProcessingStatus:
         )
 
 
+class ClassProperty(SubjectFilter):
+    subject: Literal["class_property"] = "class_property"
+    matcher: Union[AnyOf[str], AllOf[str], NoneOf[str]]
+
+    @classmethod
+    def any_of(cls, values: list[str]) -> ClassProperty:
+        return ClassProperty(
+            subject="class_property", matcher=AnyOf[str](values=values)
+        )
+
+    @classmethod
+    def all_of(cls, values: list[str]) -> ClassProperty:
+        return ClassProperty(
+            subject="class_property", matcher=AllOf[str](values=values)
+        )
+
+    @classmethod
+    def none_of(cls, values: list[str]) -> ClassProperty:
+        return ClassProperty(
+            subject="class_property", matcher=NoneOf[str](values=values)
+        )
+
+
+class ClassPropertyValue(SubjectFilter):
+    subject: Literal["class_property_value"] = "class_property_value"
+    matcher: Union[AnyOf[str], AllOf[str], NoneOf[str]]
+
+    @classmethod
+    def any_of(cls, values: list[str]) -> ClassPropertyValue:
+        return ClassPropertyValue(
+            subject="class_property_value", matcher=AnyOf[str](values=values)
+        )
+
+    @classmethod
+    def all_of(cls, values: list[str]) -> ClassPropertyValue:
+        return ClassPropertyValue(
+            subject="class_property_value", matcher=AllOf[str](values=values)
+        )
+
+    @classmethod
+    def none_of(cls, values: list[str]) -> ClassPropertyValue:
+        return ClassPropertyValue(
+            subject="class_property_value", matcher=NoneOf[str](values=values)
+        )
+
+
 class UpdatedAt(SubjectFilter):
     subject: Literal["updated_at"] = "updated_at"
     matcher: DateRange

darwin/future/data_objects/workflow.py

@@ -3,7 +3,7 @@
 from typing import List, Optional
 from uuid import UUID
 
-from pydantic import Field, model_validator
+from pydantic import ConfigDict, Field, model_validator
 
 from darwin.future.data_objects.typing import UnknownType
 from darwin.future.pydantic_base import DefaultDarwin
@@ -135,6 +135,10 @@ class WFStageConfigCore(DefaultDarwin):
     model_id: UnknownType
     threshold: UnknownType
 
+    model_config = ConfigDict(
+        protected_namespaces=()
+    )  # due to `model_type` field conflicts with a pydantic field
+
 
 class WFStageCore(DefaultDarwin):
     """

darwin/future/meta/queries/item.py

@@ -1,21 +1,22 @@
 from __future__ import annotations
 
 from functools import reduce
-from typing import Dict, Protocol
+from typing import Dict, Literal, Protocol, Union
 
 from darwin.future.core.items.archive_items import archive_list_of_items
 from darwin.future.core.items.assign_items import assign_items
 from darwin.future.core.items.delete_items import delete_list_of_items
-from darwin.future.core.items.get import list_items
+from darwin.future.core.items.get import list_items, list_items_unstable
 from darwin.future.core.items.move_items_to_folder import move_list_of_items_to_folder
 from darwin.future.core.items.restore_items import restore_list_of_items
 from darwin.future.core.items.set_item_layout import set_item_layout
 from darwin.future.core.items.set_item_priority import set_item_priority
 from darwin.future.core.items.set_stage_to_items import set_stage_to_items
 from darwin.future.core.items.tag_items import tag_items
 from darwin.future.core.items.untag_items import untag_items
-from darwin.future.core.types.common import QueryString
+from darwin.future.core.types.common import JSONDict, QueryString
 from darwin.future.core.types.query import PaginatedQuery, QueryFilter
+from darwin.future.data_objects.advanced_filters import GroupFilter, SubjectFilter
 from darwin.future.data_objects.item import ItemLayout
 from darwin.future.data_objects.sorting import SortingMethods
 from darwin.future.data_objects.workflow import WFStageCore
@@ -44,14 +45,13 @@ def _collect(self) -> Dict[int, Item]:
             else self.meta_params["dataset_id"]
         )
         team_slug = self.meta_params["team_slug"]
-        params: QueryString = reduce(
-            lambda s1, s2: s1 + s2,
-            [
-                self.page.to_query_string(),
-                *[QueryString(f.to_dict()) for f in self.filters],
-            ],
-        )
-        items_core, errors = list_items(self.client, team_slug, dataset_ids, params)
+        params = self._build_params()
+        if isinstance(params, QueryString):
+            items_core, errors = list_items(self.client, team_slug, dataset_ids, params)
+        else:
+            items_core, errors = list_items_unstable(
+                api_client=self.client, team_slug=team_slug, params=params
+            )
         offset = self.page.offset
         items = {
             i
@@ -119,6 +119,28 @@ def move_to_folder(self, path) -> None:
         filters = {"item_ids": [str(item) for item in ids]}
         move_list_of_items_to_folder(self.client, team_slug, dataset_ids, path, filters)
 
+    def _build_params(self) -> Union[QueryString, JSONDict]:
+        if self._advanced_filters is None:
+            return reduce(
+                lambda s1, s2: s1 + s2,
+                [
+                    self.page.to_query_string(),
+                    *[QueryString(f.to_dict()) for f in self.filters],
+                ],
+            )
+        if not self.meta_params["dataset_ids"] and not self.meta_params["dataset_id"]:
+            raise ValueError("Must specify dataset_ids to query items")
+        dataset_id = (
+            self.meta_params["dataset_ids"]
+            if "dataset_ids" in self.meta_params
+            else self.meta_params["dataset_id"]
+        )
+        return {
+            "dataset_ids": [dataset_id],
+            "page": self.page.model_dump(),
+            "filter": self._advanced_filters.model_dump(),
+        }
+
     def set_priority(self, priority: int) -> None:
         if "team_slug" not in self.meta_params:
             raise ValueError("Must specify team_slug to query items")
@@ -329,3 +351,36 @@ def set_stage(
         set_stage_to_items(
             self.client, team_slug, dataset_ids, stage_id, workflow_id, filters
         )
+
+    def where(
+        self,
+        *args: GroupFilter | SubjectFilter,
+        _operator: Literal["and", "or"] = "and",
+        **kwargs: str,
+    ) -> ItemQuery:
+        """Adds a filter to the query
+        This can be used with simple filters via a keyword argument
+        or with advanced filters via a GroupFilter or SubjectFilter object as args.
+
+        Args:
+            _operator (Literal["and", "or"], optional): The operator to use when combining
+
+        Raises:
+            ValueError: Raises if trying to use both simple and advanced filters
+
+        Returns:
+            ItemQuery: Self
+        """
+        if len(args) > 0 and len(kwargs) > 0:
+            raise ValueError("Cannot specify both args and kwargs")
+        if len(kwargs) > 1:
+            return super().where(**kwargs)
+        arg_list = list(args)
+        if self._advanced_filters is None:
+            self._advanced_filters = arg_list.pop(0)
+        for arg in arg_list:
+            if _operator == "and":
+                self._advanced_filters = self._advanced_filters & arg
+            if _operator == "or":
+                self._advanced_filters = self._advanced_filters | arg
+        return self