Add search filters

Author: Val Brodsky

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

@@ -0,0 +1,122 @@
+import datetime
+from enum import Enum
+from typing import List, Literal, Union
+
+from labelbox.pydantic_compat import BaseModel, Field
+
+
+class BaseSearchFilter(BaseModel):
+
+    class Config:
+        use_enum_values = True
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+        if 'operation' in res:
+            res['type'] = res.pop('operation')
+
+        # go through all the keys and convert date to string
+        for key in res:
+            if isinstance(res[key], datetime.date):
+                res[key] = res[key].isoformat()
+        return res
+
+
+class OperationType(Enum):
+    Organization = 'organization'
+    Workspace = 'workspace'
+    Tag = 'tag'
+    Stage = 'stage'
+    WorforceRequestedDate = 'workforce_requested_at'
+    WorkforceStageUpdatedDate = 'workforce_stage_updated_at'
+
+
+class IdOperator(Enum):
+    Is = 'is'
+
+
+class DateOperator(Enum):
+    Equals = 'EQUALS'
+    GreaterThanOrEqual = 'GREATER_THAN_OR_EQUAL'
+    LessThanOrEqual = 'LESS_THAN_OR_EQUAL'
+
+
+class DateRangeOperator(Enum):
+    Between = 'BETWEEN'
+
+
+class OrganizationFilter(BaseSearchFilter):
+    operation: Literal[OperationType.Organization]
+    operator: IdOperator
+    values: List[str]
+
+
+class WorkspaceFilter(BaseSearchFilter):
+    operation: Literal[OperationType.Workspace]
+    operator: IdOperator
+    values: List[str]
+
+
+class TagFilter(BaseSearchFilter):
+    operation: Literal[OperationType.Tag]
+    operator: IdOperator
+    values: List[str]
+
+
+class ProjectStageFilter(BaseSearchFilter):
+    operation: Literal[OperationType.Stage]
+    operator: IdOperator
+    values: List[str]
+
+
+class DateValue(BaseSearchFilter):
+    operator: DateOperator
+    value: datetime.date
+    # timezone: TimeZoneName = Field(default=TimeZoneName.UTC)  # type: ignore
+
+
+class WorkforceStageUpdatedFilter(BaseSearchFilter):
+    operation: Literal[OperationType.WorkforceStageUpdatedDate]
+    value: DateValue
+
+
+class WorkforceRequestedDateFilter(BaseSearchFilter):
+    operation: Literal[OperationType.WorforceRequestedDate]
+    value: DateValue
+
+
+class DateRange(BaseSearchFilter):
+    min: datetime.date
+    max: datetime.date
+
+
+class DateRangeValue(BaseSearchFilter):
+    operator: DateRangeOperator
+    value: DateRange
+
+
+class WorkforceRequestedDateRangeFilter(BaseSearchFilter):
+    operation: Literal[OperationType.WorforceRequestedDate]
+    # timezone: TimeZoneName = Field(default=TimeZoneName.UTC)  # type: ignore
+    value: DateRangeValue
+
+
+class WorkforceStageUpdatedRangeFilter(BaseSearchFilter):
+    operation: Literal[OperationType.WorkforceStageUpdatedDate]
+    # timezone: TimeZoneName = Field(default=TimeZoneName.UTC)  # type: ignore
+    value: DateRangeValue
+
+
+SearchFilters = Union[OrganizationFilter, WorkspaceFilter, TagFilter,
+                      ProjectStageFilter, WorkforceRequestedDateFilter,
+                      WorkforceStageUpdatedFilter,
+                      WorkforceRequestedDateRangeFilter,
+                      WorkforceStageUpdatedRangeFilter]
+
+
+def _build_search_filter(filter: List[SearchFilters]):
+    operation_types = {f.operation for f in filter}
+    if len(operation_types) < len(filter):
+        raise ValueError("Only one filter per operation type is allowed")
+
+    return [f.dict() for f in filter]

libs/labelbox/tests/unit/test_unit_search_filters.py

@@ -0,0 +1,137 @@
+import pytest
+
+from labelbox.schema.search_filters import DateOperator, DateRange, DateRangeOperator, DateRangeValue, DateValue, IdOperator, OperationType, OrganizationFilter, ProjectStageFilter, TagFilter, WorkforceRequestedDateFilter, WorkforceRequestedDateRangeFilter, WorkforceStageUpdatedFilter, WorkforceStageUpdatedRangeFilter, WorkspaceFilter, _build_search_filter
+from test.test_pdb import pdb
+
+
+def test_duplicate_filters():
+    filters = [
+        OrganizationFilter(operation=OperationType.Organization,
+                           operator=IdOperator.Is,
+                           values=["clphb4vd7000cd2wv1ktu5cwa"]),
+        WorkspaceFilter(operation=OperationType.Workspace,
+                        operator=IdOperator.Is,
+                        values=["clphb4vd7000cd2wv1ktu5cwa"]),
+        TagFilter(operation=OperationType.Tag,
+                  operator=IdOperator.Is,
+                  values=["tag"]),
+        ProjectStageFilter(operation=OperationType.Stage,
+                           operator=IdOperator.Is,
+                           values=["done"]),
+        WorkforceRequestedDateFilter(
+            operation=OperationType.WorforceRequestedDate,
+            operator=DateOperator.Equals,
+            value="2024-01-01"),
+        WorkforceStageUpdatedFilter(
+            operation=OperationType.WorkforceStageUpdatedDate,
+            operator=DateOperator.Equals,
+            value="2024-01-01"),
+        WorkforceRequestedDateRangeFilter(
+            operation=OperationType.WorforceRequestedDate,
+            value=DateRangeValue(operator=DateRangeOperator.Between,
+                                 value=DateRange(min="2024-01-01",
+                                                 max="2025-01-01"))),
+        WorkforceStageUpdatedRangeFilter(
+            operation=OperationType.WorkforceStageUpdatedDate,
+            value=DateRangeValue(operator=DateRangeOperator.Between,
+                                 value=DateRange(min="2024-01-01",
+                                                 max="2025-01-01")))
+    ]
+
+    with pytest.raises(ValueError):
+        _build_search_filter(filters)
+
+
+def test_id_filters():
+    filters = [
+        OrganizationFilter(operation=OperationType.Organization,
+                           operator=IdOperator.Is,
+                           values=["clphb4vd7000cd2wv1ktu5cwa"]),
+        WorkspaceFilter(operation=OperationType.Workspace,
+                        operator=IdOperator.Is,
+                        values=["clphb4vd7000cd2wv1ktu5cwa"]),
+        TagFilter(operation=OperationType.Tag,
+                  operator=IdOperator.Is,
+                  values=["tag"]),
+        ProjectStageFilter(operation=OperationType.Stage,
+                           operator=IdOperator.Is,
+                           values=["requested"]),
+    ]
+
+    assert _build_search_filter(filters) == [{
+        "operator": "is",
+        "values": ["clphb4vd7000cd2wv1ktu5cwa"],
+        "type": "organization"
+    }, {
+        "operator": "is",
+        "values": ["clphb4vd7000cd2wv1ktu5cwa"],
+        "type": "workspace"
+    }, {
+        "operator": "is",
+        "values": ["tag"],
+        "type": "tag"
+    }, {
+        "operator": "is",
+        "values": ["requested"],
+        "type": "stage"
+    }]
+
+
+def test_date_filters():
+    filters = [
+        WorkforceRequestedDateFilter(
+            operation=OperationType.WorforceRequestedDate,
+            value=DateValue(operator=DateOperator.GreaterThanOrEqual,
+                            value="2024-01-01")),
+        WorkforceStageUpdatedFilter(
+            operation=OperationType.WorkforceStageUpdatedDate,
+            value=DateValue(operator=DateOperator.LessThanOrEqual,
+                            value="2025-01-01")),
+    ]
+    assert _build_search_filter(filters) == [{
+        "type": "workforce_requested_at",
+        "value": {
+            "operator": "GREATER_THAN_OR_EQUAL",
+            "value": "2024-01-01",
+        }
+    }, {
+        "type": "workforce_stage_updated_at",
+        "value": {
+            "operator": "LESS_THAN_OR_EQUAL",
+            "value": "2025-01-01",
+        }
+    }]
+
+
+def test_date_range_filters():
+    filters = [
+        WorkforceRequestedDateRangeFilter(
+            operation=OperationType.WorforceRequestedDate,
+            value=DateRangeValue(operator=DateRangeOperator.Between,
+                                 value=DateRange(min="2024-01-01",
+                                                 max="2025-01-01"))),
+        WorkforceStageUpdatedRangeFilter(
+            operation=OperationType.WorkforceStageUpdatedDate,
+            value=DateRangeValue(operator=DateRangeOperator.Between,
+                                 value=DateRange(min="2024-01-01",
+                                                 max="2025-01-01")))
+    ]
+    assert _build_search_filter(filters) == [{
+        "value": {
+            "operator": "BETWEEN",
+            "value": {
+                "min": "2024-01-01",
+                "max": "2025-01-01"
+            }
+        },
+        "type": "workforce_requested_at"
+    }, {
+        "value": {
+            "operator": "BETWEEN",
+            "value": {
+                "min": "2024-01-01",
+                "max": "2025-01-01"
+            }
+        },
+        "type": "workforce_stage_updated_at"
+    }]