[PLT-1412] Vb/add task filters plt 1412 (#1773)

vbrodsky · web-flow · commit 4171bec8f71c · 2024-08-21T11:29:10.000-07:00
diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py
@@ -2435,7 +2435,7 @@ def get_labeling_service_dashboards(
             >>> seven_days_ago = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d")
             >>> workforce_requested_filter_before = WorkforceRequestedDateFilter(
             >>>     operation=OperationType.WorforceRequestedDate,
-            >>>     value=DateValue(operator=DateOperator.GreaterThanOrEqual,
+            >>>     value=DateValue(operator=RangeOperatorWithSingleValue.GreaterThanOrEqual,
             >>>                     value=seven_days_ago))
             >>> labeling_service_dashboard = [ld for ld in project.client.get_labeling_service_dashboards(search_query=[workforce_requested_filter_before])]
 
diff --git a/libs/labelbox/src/labelbox/schema/search_filters.py b/libs/labelbox/src/labelbox/schema/search_filters.py
@@ -36,6 +36,8 @@ class OperationType(Enum):
     Stage = 'stage'
     WorforceRequestedDate = 'workforce_requested_at'
     WorkforceStageUpdatedDate = 'workforce_stage_updated_at'
+    TaskCompletedCount = 'task_completed_count'
+    TaskRemainingCount = 'task_remaining_count'
 
 
 class IdOperator(Enum):
@@ -45,7 +47,7 @@ class IdOperator(Enum):
     Is = 'is'
 
 
-class DateOperator(Enum):
+class RangeOperatorWithSingleValue(Enum):
     """
     Supported operators for dates
     """
@@ -54,7 +56,7 @@ class DateOperator(Enum):
     LessThanOrEqual = 'LESS_THAN_OR_EQUAL'
 
 
-class DateRangeOperator(Enum):
+class RangeOperatorWithValue(Enum):
     """
     Supported operators for date ranges
     """
@@ -65,7 +67,7 @@ class OrganizationFilter(BaseSearchFilter):
     """
     Filter for organization
     """
-    operation: Literal[OperationType.Organization]
+    operation: Literal[OperationType.Organization] = OperationType.Organization
     operator: IdOperator
     values: List[str]
 
@@ -74,7 +76,7 @@ class WorkspaceFilter(BaseSearchFilter):
     """
     Filter for workspace
     """
-    operation: Literal[OperationType.Workspace]
+    operation: Literal[OperationType.Workspace] = OperationType.Workspace
     operator: IdOperator
     values: List[str]
 
@@ -83,7 +85,7 @@ class TagFilter(BaseSearchFilter):
     """
     Filter for project tags
     """
-    operation: Literal[OperationType.Tag]
+    operation: Literal[OperationType.Tag] = OperationType.Tag
     operator: IdOperator
     values: List[str]
 
@@ -92,7 +94,7 @@ class ProjectStageFilter(BaseSearchFilter):
     """
     Filter labelbox service / aka project stages
     """
-    operation: Literal[OperationType.Stage]
+    operation: Literal[OperationType.Stage] = OperationType.Stage
     operator: IdOperator
     values: List[str]
 
@@ -109,23 +111,32 @@ class DateValue(BaseSearchFilter):
             so for a string '2024-01-01' that is run on a computer in PST, we would convert it to '2024-01-01T08:00:00Z'
             while the same string in EST will get converted to '2024-01-01T05:00:00Z'
     """
-    operator: DateOperator
+    operator: RangeOperatorWithSingleValue
     value: datetime.datetime
 
 
+class IntegerValue(BaseSearchFilter):
+    operator: RangeOperatorWithSingleValue
+    value: int
+
+
 class WorkforceStageUpdatedFilter(BaseSearchFilter):
     """
     Filter for workforce stage updated date
     """
-    operation: Literal[OperationType.WorkforceStageUpdatedDate]
+    operation: Literal[
+        OperationType.
+        WorkforceStageUpdatedDate] = OperationType.WorkforceStageUpdatedDate
     value: DateValue
 
 
 class WorkforceRequestedDateFilter(BaseSearchFilter):
     """
     Filter for workforce requested date
     """
-    operation: Literal[OperationType.WorforceRequestedDate]
+    operation: Literal[
+        OperationType.
+        WorforceRequestedDate] = OperationType.WorforceRequestedDate
     value: DateValue
 
 
@@ -141,31 +152,54 @@ class DateRangeValue(BaseSearchFilter):
     """
     Date range value for a search filter 
     """
-    operator: DateRangeOperator
+    operator: RangeOperatorWithValue
     value: DateRange
 
 
 class WorkforceRequestedDateRangeFilter(BaseSearchFilter):
     """
     Filter for workforce requested date range
     """
-    operation: Literal[OperationType.WorforceRequestedDate]
+    operation: Literal[
+        OperationType.
+        WorforceRequestedDate] = OperationType.WorforceRequestedDate
     value: DateRangeValue
 
 
 class WorkforceStageUpdatedRangeFilter(BaseSearchFilter):
     """
     Filter for workforce stage updated date range
     """
-    operation: Literal[OperationType.WorkforceStageUpdatedDate]
+    operation: Literal[
+        OperationType.
+        WorkforceStageUpdatedDate] = OperationType.WorkforceStageUpdatedDate
     value: DateRangeValue
 
 
+class TaskCompletedCountFilter(BaseSearchFilter):
+    """
+    Filter for completed tasks count
+    """
+    operation: Literal[
+        OperationType.TaskCompletedCount] = OperationType.TaskCompletedCount
+    value: IntegerValue
+
+
+class TaskRemainingCountFilter(BaseSearchFilter):
+    """
+    Filter for remaining tasks count
+    """
+    operation: Literal[
+        OperationType.TaskRemainingCount] = OperationType.TaskRemainingCount
+    value: IntegerValue
+
+
 SearchFilter = Union[OrganizationFilter, WorkspaceFilter, TagFilter,
                      ProjectStageFilter, WorkforceRequestedDateFilter,
                      WorkforceStageUpdatedFilter,
                      WorkforceRequestedDateRangeFilter,
-                     WorkforceStageUpdatedRangeFilter]
+                     WorkforceStageUpdatedRangeFilter, TaskCompletedCountFilter,
+                     TaskRemainingCountFilter]
 
 
 def _dict_to_graphql_string(d: Union[dict, list, str, int]) -> str:
diff --git a/libs/labelbox/tests/integration/test_labeling_dashboard.py b/libs/labelbox/tests/integration/test_labeling_dashboard.py
@@ -1,8 +1,8 @@
 from datetime import datetime, timedelta
 from labelbox.schema.labeling_service import LabelingServiceStatus
-from labelbox.schema.search_filters import DateOperator, DateRange, DateRangeOperator, DateRangeValue, DateValue, IdOperator, OperationType, OrganizationFilter, WorkforceRequestedDateFilter, WorkforceRequestedDateRangeFilter, WorkspaceFilter
 from labelbox.schema.ontology_kind import EditorTaskType
 from labelbox.schema.media_type import MediaType
+from labelbox.schema.search_filters import IntegerValue, RangeOperatorWithSingleValue, DateRange, RangeOperatorWithValue, DateRangeValue, DateValue, IdOperator, OperationType, OrganizationFilter, TaskCompletedCountFilter, WorkforceRequestedDateFilter, WorkforceRequestedDateRangeFilter, WorkspaceFilter, TaskRemainingCountFilter
 
 
 def test_request_labeling_service_dashboard(rand_gen,
@@ -48,12 +48,13 @@ def test_request_labeling_service_dashboard_filters(requested_labeling_service):
 
     workforce_requested_filter_before = WorkforceRequestedDateFilter(
         operation=OperationType.WorforceRequestedDate,
-        value=DateValue(operator=DateOperator.GreaterThanOrEqual,
-                        value=datetime.strptime("2024-01-01", "%Y-%m-%d")))
+        value=DateValue(
+            operator=RangeOperatorWithSingleValue.GreaterThanOrEqual,
+            value=datetime.strptime("2024-01-01", "%Y-%m-%d")))
     year_from_now = (datetime.now() + timedelta(days=365))
     workforce_requested_filter_after = WorkforceRequestedDateFilter(
         operation=OperationType.WorforceRequestedDate,
-        value=DateValue(operator=DateOperator.LessThanOrEqual,
+        value=DateValue(operator=RangeOperatorWithSingleValue.LessThanOrEqual,
                         value=year_from_now))
 
     labeling_service_dashboard = [
@@ -66,7 +67,7 @@ def test_request_labeling_service_dashboard_filters(requested_labeling_service):
 
     workforce_date_range_filter = WorkforceRequestedDateRangeFilter(
         operation=OperationType.WorforceRequestedDate,
-        value=DateRangeValue(operator=DateRangeOperator.Between,
+        value=DateRangeValue(operator=RangeOperatorWithValue.Between,
                              value=DateRange(min="2024-01-01T00:00:00-0800",
                                              max=year_from_now)))
 
@@ -90,3 +91,18 @@ def test_request_labeling_service_dashboard_filters(requested_labeling_service):
     labeling_service_dashboard = project.client.get_labeling_service_dashboards(
     ).get_one()
     assert labeling_service_dashboard
+
+    task_done_count_filter = TaskCompletedCountFilter(
+        operation=OperationType.TaskCompletedCount,
+        value=IntegerValue(
+            operator=RangeOperatorWithSingleValue.GreaterThanOrEqual, value=0))
+    task_remaining_count_filter = TaskRemainingCountFilter(
+        operation=OperationType.TaskRemainingCount,
+        value=IntegerValue(
+            operator=RangeOperatorWithSingleValue.GreaterThanOrEqual, value=0))
+
+    labeling_service_dashboard = [
+        ld for ld in project.client.get_labeling_service_dashboards(
+            search_query=[task_done_count_filter, task_remaining_count_filter])
+    ][0]
+    assert labeling_service_dashboard is not None
diff --git a/libs/labelbox/tests/unit/test_unit_search_filters.py b/libs/labelbox/tests/unit/test_unit_search_filters.py
@@ -1,22 +1,16 @@
 from datetime import datetime
-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 labelbox.schema.search_filters import IntegerValue, RangeOperatorWithSingleValue, DateRange, RangeOperatorWithValue, DateRangeValue, DateValue, IdOperator, OperationType, OrganizationFilter, ProjectStageFilter, TagFilter, TaskCompletedCountFilter, TaskRemainingCountFilter, WorkforceRequestedDateFilter, WorkforceRequestedDateRangeFilter, WorkforceStageUpdatedFilter, WorkforceStageUpdatedRangeFilter, WorkspaceFilter, build_search_filter
 from labelbox.utils import format_iso_datetime
 
 
 def test_id_filters():
     filters = [
-        OrganizationFilter(operation=OperationType.Organization,
-                           operator=IdOperator.Is,
+        OrganizationFilter(operator=IdOperator.Is,
                            values=["clphb4vd7000cd2wv1ktu5cwa"]),
-        WorkspaceFilter(operation=OperationType.Workspace,
-                        operator=IdOperator.Is,
+        WorkspaceFilter(operator=IdOperator.Is,
                         values=["clphb4vd7000cd2wv1ktu5cwa"]),
-        TagFilter(operation=OperationType.Tag,
-                  operator=IdOperator.Is,
-                  values=["tag"]),
-        ProjectStageFilter(operation=OperationType.Stage,
-                           operator=IdOperator.Is,
-                           values=["requested"]),
+        TagFilter(operator=IdOperator.Is, values=["tag"]),
+        ProjectStageFilter(operator=IdOperator.Is, values=["requested"]),
     ]
 
     assert build_search_filter(
@@ -29,14 +23,12 @@ def test_date_filters():
     local_time_end = datetime.strptime("2025-01-01", "%Y-%m-%d")
 
     filters = [
-        WorkforceRequestedDateFilter(
-            operation=OperationType.WorforceRequestedDate,
-            value=DateValue(operator=DateOperator.GreaterThanOrEqual,
-                            value=local_time_start)),
-        WorkforceStageUpdatedFilter(
-            operation=OperationType.WorkforceStageUpdatedDate,
-            value=DateValue(operator=DateOperator.LessThanOrEqual,
-                            value=local_time_end)),
+        WorkforceRequestedDateFilter(value=DateValue(
+            operator=RangeOperatorWithSingleValue.GreaterThanOrEqual,
+            value=local_time_start)),
+        WorkforceStageUpdatedFilter(value=DateValue(
+            operator=RangeOperatorWithSingleValue.LessThanOrEqual,
+            value=local_time_end)),
     ]
     expected_start = format_iso_datetime(local_time_start)
     expected_end = format_iso_datetime(local_time_end)
@@ -47,25 +39,31 @@ def test_date_filters():
 
 def test_date_range_filters():
     filters = [
-        WorkforceRequestedDateRangeFilter(
-            operation=OperationType.WorforceRequestedDate,
-            value=DateRangeValue(operator=DateRangeOperator.Between,
-                                 value=DateRange(min=datetime.strptime(
-                                     "2024-01-01T00:00:00-0800",
-                                     "%Y-%m-%dT%H:%M:%S%z"),
-                                                 max=datetime.strptime(
-                                                     "2025-01-01T00:00:00-0800",
-                                                     "%Y-%m-%dT%H:%M:%S%z")))),
-        WorkforceStageUpdatedRangeFilter(
-            operation=OperationType.WorkforceStageUpdatedDate,
-            value=DateRangeValue(operator=DateRangeOperator.Between,
-                                 value=DateRange(min=datetime.strptime(
-                                     "2024-01-01T00:00:00-0800",
-                                     "%Y-%m-%dT%H:%M:%S%z"),
-                                                 max=datetime.strptime(
-                                                     "2025-01-01T00:00:00-0800",
-                                                     "%Y-%m-%dT%H:%M:%S%z")))),
+        WorkforceRequestedDateRangeFilter(value=DateRangeValue(
+            operator=RangeOperatorWithValue.Between,
+            value=DateRange(min=datetime.strptime("2024-01-01T00:00:00-0800",
+                                                  "%Y-%m-%dT%H:%M:%S%z"),
+                            max=datetime.strptime("2025-01-01T00:00:00-0800",
+                                                  "%Y-%m-%dT%H:%M:%S%z")))),
+        WorkforceStageUpdatedRangeFilter(value=DateRangeValue(
+            operator=RangeOperatorWithValue.Between,
+            value=DateRange(min=datetime.strptime("2024-01-01T00:00:00-0800",
+                                                  "%Y-%m-%dT%H:%M:%S%z"),
+                            max=datetime.strptime("2025-01-01T00:00:00-0800",
+                                                  "%Y-%m-%dT%H:%M:%S%z")))),
     ]
     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"}]'
+
+
+def test_task_count_filters():
+    filters = [
+        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"}]'
+    assert build_search_filter(filters) == expected
