cleanup

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

packages/pynumaflow/pynumaflow/accumulator/_dtypes.py

@@ -17,9 +17,9 @@ class WindowOperation(IntEnum):
     Enumerate the type of Window operation received.
     """
 
-    OPEN = (0,)
-    CLOSE = (1,)
-    APPEND = (2,)
+    OPEN = 0
+    CLOSE = 1
+    APPEND = 2
 
 
 @dataclass(init=False, slots=True)

packages/pynumaflow/pynumaflow/accumulator/servicer/task_manager.py

@@ -13,7 +13,7 @@
     Datum,
     _AccumulatorBuilderClass,
     AccumulatorAsyncCallable,
-    WindowOperation,
+    WindowOperation, AccumulatorRequest,
 )
 from pynumaflow.proto.accumulator import accumulator_pb2
 from pynumaflow.shared.asynciter import NonBlockingIterator
@@ -93,7 +93,7 @@ async def stream_send_eof(self):
         for unified_key in task_keys:
             await self.tasks[unified_key].iterator.put(STREAM_EOF)
 
-    async def close_task(self, req):
+    async def close_task(self, req: AccumulatorRequest):
         """
         Closes a running accumulator task for a given key.
         Based on the request we compute the unique key, and then
@@ -121,7 +121,7 @@ async def close_task(self, req):
             # Put the exception in the result queue
             await self.global_result_queue.put(err)
 
-    async def create_task(self, req):
+    async def create_task(self, req: AccumulatorRequest):
         """
         Creates a new accumulator task for the given request.
         Based on the request we compute a unique key, and then
@@ -141,7 +141,7 @@ async def create_task(self, req):
             # Create a new result queue for the current task
             # We create a new result queue for each task, so that
             # the results of the accumulator operation can be sent to the
-            # the global result queue, which in turn sends the results
+            # global result queue, which in turn sends the results
             # to the client.
             res_queue = NonBlockingIterator()
 
@@ -175,7 +175,7 @@ async def create_task(self, req):
         # Put the request in the iterator
         await curr_task.iterator.put(d)
 
-    async def send_datum_to_task(self, req):
+    async def send_datum_to_task(self, req: AccumulatorRequest):
         """
         Appends the request to the existing window reduce task.
         If the task does not exist, create it.
@@ -220,7 +220,7 @@ async def __invoke_accumulator(
             await self.global_result_queue.put(err)
 
     async def process_input_stream(
-        self, request_iterator: AsyncIterable[accumulator_pb2.AccumulatorRequest]
+        self, request_iterator: AsyncIterable[AccumulatorRequest]
     ):
         # Start iterating through the request iterator and create tasks
         # based on the operation type received.
@@ -230,15 +230,15 @@ async def process_input_stream(
                 request_count += 1
                 # check whether the request is an open, append, or close operation
                 match request.operation:
-                    case int(WindowOperation.OPEN):
+                    case WindowOperation.OPEN:
                         # create a new task for the open operation and
                         # put the request in the task iterator
                         await self.create_task(request)
-                    case int(WindowOperation.APPEND):
+                    case WindowOperation.APPEND:
                         # append the task data to the existing task
                         # if the task does not exist, create a new task
                         await self.send_datum_to_task(request)
-                    case int(WindowOperation.CLOSE):
+                    case WindowOperation.CLOSE:
                         # close the current task for req
                         await self.close_task(request)
                     case _: