Fix search filters

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/search_filters.py

@@ -1,30 +1,24 @@
 import datetime
 from enum import Enum
-from typing import List, Literal, Union
+from typing import List, Union
+from pydantic import PlainSerializer, BaseModel, Field
+
+from typing_extensions import Annotated
 
-from pydantic import BaseModel, Field
 from labelbox.schema.labeling_service_status import LabelingServiceStatus
 from labelbox.utils import format_iso_datetime
-from pydantic.config import ConfigDict
 
 
 class BaseSearchFilter(BaseModel):
     """
     Shared code for all search filters
     """
 
-    model_config = ConfigDict(use_enum_values=True)
-
-    def dict(self, *args, **kwargs):
-        res = super().dict(*args, **kwargs)
-        # go through all the keys and convert date to string
-        for key in res:
-            if isinstance(res[key], datetime.datetime):
-                res[key] = format_iso_datetime(res[key])
-        return res
+    class Config:
+        use_enum_values = True
 
 
-class OperationType(Enum):
+class OperationTypeEnum(Enum):
     """
     Supported search entity types
     Each type corresponds to a different filter class
@@ -40,6 +34,13 @@ class OperationType(Enum):
     TaskRemainingCount = 'task_remaining_count'
 
 
+OperationType = Annotated[OperationTypeEnum,
+                          PlainSerializer(lambda x: x.value, return_type=str)]
+
+IsoDatetimeType = Annotated[datetime.datetime,
+                            PlainSerializer(format_iso_datetime)]
+
+
 class IdOperator(Enum):
     """
     Supported operators for ids like org ids, workspace ids, etc
@@ -75,7 +76,8 @@ class OrganizationFilter(BaseSearchFilter):
     """
     Filter for organization to which projects belong
     """
-    operation: Literal[OperationType.Organization] = OperationType.Organization
+    operation: OperationType = Field(default=OperationTypeEnum.Organization,
+                                     serialization_alias='type')
     operator: IdOperator
     values: List[str]
 
@@ -84,9 +86,10 @@ class SharedWithOrganizationFilter(BaseSearchFilter):
     """
     Find project shared with the organization (i.e. not having this organization as a tenantId)
     """
-    operation: Literal[
-        OperationType.
-        SharedWithOrganization] = OperationType.SharedWithOrganization
+
+    operation: OperationType = Field(
+        default=OperationTypeEnum.SharedWithOrganization,
+        serialization_alias='type')
     operator: IdOperator
     values: List[str]
 
@@ -95,7 +98,8 @@ class WorkspaceFilter(BaseSearchFilter):
     """
     Filter for workspace
     """
-    operation: Literal[OperationType.Workspace] = OperationType.Workspace
+    operation: OperationType = Field(default=OperationTypeEnum.Workspace,
+                                     serialization_alias='type')
     operator: IdOperator
     values: List[str]
 
@@ -104,7 +108,8 @@ class TagFilter(BaseSearchFilter):
     """
     Filter for project tags
     """
-    operation: Literal[OperationType.Tag] = OperationType.Tag
+    operation: OperationType = Field(default=OperationTypeEnum.Tag,
+                                     serialization_alias='type')
     operator: IdOperator
     values: List[str]
 
@@ -114,7 +119,8 @@ class ProjectStageFilter(BaseSearchFilter):
     Filter labelbox service / aka project stages
         Stages are: requested, in_progress, completed etc. as described by LabelingServiceStatus
     """
-    operation: Literal[OperationType.Stage] = OperationType.Stage
+    operation: OperationType = Field(default=OperationTypeEnum.Stage,
+                                     serialization_alias='type')
     operator: IdOperator
     values: List[LabelingServiceStatus]
 
@@ -132,7 +138,7 @@ class DateValue(BaseSearchFilter):
             while the same string in EST will get converted to '2024-01-01T05:00:00Z'
     """
     operator: RangeDateTimeOperatorWithSingleValue
-    value: datetime.datetime
+    value: IsoDatetimeType
 
 
 class IntegerValue(BaseSearchFilter):
@@ -144,28 +150,28 @@ class WorkforceStageUpdatedFilter(BaseSearchFilter):
     """
     Filter for workforce stage updated date
     """
-    operation: Literal[
-        OperationType.
-        WorkforceStageUpdatedDate] = OperationType.WorkforceStageUpdatedDate
+    operation: OperationType = Field(
+        default=OperationTypeEnum.WorkforceStageUpdatedDate,
+        serialization_alias='type')
     value: DateValue
 
 
 class WorkforceRequestedDateFilter(BaseSearchFilter):
     """
     Filter for workforce requested date
     """
-    operation: Literal[
-        OperationType.
-        WorforceRequestedDate] = OperationType.WorforceRequestedDate
+    operation: OperationType = Field(
+        default=OperationTypeEnum.WorforceRequestedDate,
+        serialization_alias='type')
     value: DateValue
 
 
 class DateRange(BaseSearchFilter):
     """
     Date range for a search filter
     """
-    min: datetime.datetime
-    max: datetime.datetime
+    min: IsoDatetimeType
+    max: IsoDatetimeType
 
 
 class DateRangeValue(BaseSearchFilter):
@@ -180,19 +186,19 @@ class WorkforceRequestedDateRangeFilter(BaseSearchFilter):
     """
     Filter for workforce requested date range
     """
-    operation: Literal[
-        OperationType.
-        WorforceRequestedDate] = OperationType.WorforceRequestedDate
+    operation: OperationType = Field(
+        default=OperationTypeEnum.WorforceRequestedDate,
+        serialization_alias='type')
     value: DateRangeValue
 
 
 class WorkforceStageUpdatedRangeFilter(BaseSearchFilter):
     """
     Filter for workforce stage updated date range
     """
-    operation: Literal[
-        OperationType.
-        WorkforceStageUpdatedDate] = OperationType.WorkforceStageUpdatedDate
+    operation: OperationType = Field(
+        default=OperationTypeEnum.WorkforceStageUpdatedDate,
+        serialization_alias='type')
     value: DateRangeValue
 
 
@@ -201,20 +207,19 @@ class TaskCompletedCountFilter(BaseSearchFilter):
     Filter for completed tasks count
         A task maps to a data row. Task completed should map to a data row in a labeling queue DONE
     """
-    operation: Literal[
-        OperationType.TaskCompletedCount] = Field(default=OperationType.TaskCompletedCount, serialization_alias='type')
+    operation: OperationType = Field(
+        default=OperationTypeEnum.TaskCompletedCount,
+        serialization_alias='type')
     value: IntegerValue
 
 
-
-
-
 class TaskRemainingCountFilter(BaseSearchFilter):
     """
     Filter for remaining tasks count. Reverse of TaskCompletedCountFilter
     """
-    operation: Literal[
-        OperationType.TaskRemainingCount] = Field(OperationType.TaskRemainingCount, serialization_alias='type')
+    operation: OperationType = Field(
+        default=OperationTypeEnum.TaskRemainingCount,
+        serialization_alias='type')
     value: IntegerValue
 
 
@@ -242,5 +247,7 @@ def build_search_filter(filter: List[SearchFilter]):
     """
     Converts a list of search filters to a graphql string
     """
-    filters = [_dict_to_graphql_string(f.model_dump(by_alias=True)) for f in filter]
+    filters = [
+        _dict_to_graphql_string(f.model_dump(by_alias=True)) for f in filter
+    ]
     return "[" + ", ".join(filters) + "]"

libs/labelbox/tests/unit/test_unit_search_filters.py

@@ -19,7 +19,7 @@ def test_id_filters():
 
     assert build_search_filter(
         filters
-    ) == '[{operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"], type: "organization_id"}, {operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"], type: "shared_with_organizations"}, {operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"], type: "workspace"}, {operator: "is", values: ["tag"], type: "tag"}, {operator: "is", values: ["REQUESTED"], type: "stage"}]'
+    ) == '[{type: "organization_id", operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"]}, {type: "shared_with_organizations", operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"]}, {type: "workspace", operator: "is", values: ["clphb4vd7000cd2wv1ktu5cwa"]}, {type: "tag", operator: "is", values: ["tag"]}, {type: "stage", operator: "is", values: ["REQUESTED"]}]'
 
 
 def test_date_filters():
@@ -37,7 +37,7 @@ def test_date_filters():
     expected_start = format_iso_datetime(local_time_start)
     expected_end = format_iso_datetime(local_time_end)
 
-    expected = '[{value: {operator: "GREATER_THAN_OR_EQUAL", value: "' + expected_start + '"}, type: "workforce_requested_at"}, {value: {operator: "LESS_THAN_OR_EQUAL", value: "' + expected_end + '"}, type: "workforce_stage_updated_at"}]'
+    expected = '[{type: "workforce_requested_at", value: {operator: "GREATER_THAN_OR_EQUAL", value: "' + expected_start + '"}}, {type: "workforce_stage_updated_at", value: {operator: "LESS_THAN_OR_EQUAL", value: "' + expected_end + '"}}]'
     assert build_search_filter(filters) == expected
 
 
@@ -58,16 +58,16 @@ def test_date_range_filters():
     ]
     assert build_search_filter(
         filters
-    ) == '[{value: {operator: "BETWEEN", value: {min: "2024-01-01T08:00:00Z", max: "2025-01-01T08:00:00Z"}}, type: "workforce_requested_at"}, {value: {operator: "BETWEEN", value: {min: "2024-01-01T08:00:00Z", max: "2025-01-01T08:00:00Z"}}, type: "workforce_stage_updated_at"}]'
+    ) == '[{type: "workforce_requested_at", value: {operator: "BETWEEN", value: {min: "2024-01-01T08:00:00Z", max: "2025-01-01T08:00:00Z"}}}, {type: "workforce_stage_updated_at", value: {operator: "BETWEEN", value: {min: "2024-01-01T08:00:00Z", max: "2025-01-01T08:00:00Z"}}}]'
 
 
 def test_task_count_filters():
     filters = [
-        TaskCompletedCountFilter(value=IntegerValue(operator=RangeOperatorWithSingleValue.GreaterThanOrEqual, value=1)),
-        # TaskRemainingCountFilter(value=IntegerValue(
-        #     operator=RangeOperatorWithSingleValue.LessThanOrEqual, value=10)),
+        TaskCompletedCountFilter(value=IntegerValue(
+            operator=RangeOperatorWithSingleValue.GreaterThanOrEqual, value=1)),
+        TaskRemainingCountFilter(value=IntegerValue(
+            operator=RangeOperatorWithSingleValue.LessThanOrEqual, value=10)),
     ]
 
-    # expected = '[{value: {operator: "GREATER_THAN_OR_EQUAL", value: 1}, type: "task_completed_count"}, {value: {operator: "LESS_THAN_OR_EQUAL", value: 10}, type: "task_remaining_count"}]'
-    expected = '[{value: {operator: "GREATER_THAN_OR_EQUAL", value: 1}, type: "task_completed_count"}, {value: {operator: "LESS_THAN_OR_EQUAL", value: 10}, type: "task_remaining_count"}]'
+    expected = '[{type: "task_completed_count", value: {operator: "GREATER_THAN_OR_EQUAL", value: 1}}, {type: "task_remaining_count", value: {operator: "LESS_THAN_OR_EQUAL", value: 10}}]'
     assert build_search_filter(filters) == expected