fix: lint

Signed-off-by: srao12 <Shrivardhan_Rao@intuit.com>

Author: srao12

examples/accumulator/streamsorter/example.py

@@ -2,7 +2,6 @@
 import os
 from collections.abc import AsyncIterable
 from datetime import datetime
-from typing import List
 
 from pynumaflow import setup_logging
 from pynumaflow.accumulator import Accumulator, AccumulatorAsyncServer
@@ -21,7 +20,7 @@ class StreamSorter(Accumulator):
     def __init__(self):
         _LOGGER.info("StreamSorter initialized")
         self.latest_wm = datetime.fromtimestamp(-1)
-        self.sorted_buffer: List[Datum] = []
+        self.sorted_buffer: list[Datum] = []
 
     async def handler(
         self,

pynumaflow/accumulator/_dtypes.py

@@ -362,10 +362,10 @@ def id(self) -> str:
     @classmethod
     def from_datum(cls, datum: Datum):
         """Create a Message instance from a Datum object.
-        
+
         Args:
             datum: The Datum object to convert
-            
+
         Returns:
             Message: A new Message instance with data from the datum
         """
@@ -375,7 +375,7 @@ def from_datum(cls, datum: Datum):
             watermark=datum.watermark,
             event_time=datum.event_time,
             headers=datum.headers,
-            id=datum.id
+            id=datum.id,
         )
 
 

pynumaflow/accumulator/servicer/async_servicer.py

@@ -1,7 +1,6 @@
 import asyncio
 from collections.abc import AsyncIterable
 from typing import Union
-import logging
 
 from google.protobuf import empty_pb2 as _empty_pb2
 
@@ -31,7 +30,7 @@ async def datum_generator(
             slot=d.operation.keyedWindow.slot,
             keys=list(d.operation.keyedWindow.keys),
         )
-        
+
         accumulator_request = AccumulatorRequest(
             operation=d.operation.event,
             keyed_window=keyed_window,  # Use the new parameter name
@@ -105,23 +104,16 @@ async def AccumulateFn(
         # This is a special message that indicates the end of the processing for a window
         # When we get this message, we send an EOF message to the client
         try:
-            logging.info("[ACCUMULATOR_DEBUG] Starting to read from consumer")
             async for msg in consumer:
-                logging.info(f"[ACCUMULATOR_DEBUG] Received message type: {type(msg)}")
                 # If the message is an exception, we raise the exception
                 if isinstance(msg, BaseException):
-                    logging.info(f"[ACCUMULATOR_DEBUG] Found exception: {msg}")
-                    logging.info(f"[ACCUMULATOR_DEBUG] Calling handle_async_error with exception: {repr(msg)}")
                     await handle_async_error(context, msg, ERR_UDF_EXCEPTION_STRING)
-                    logging.info(f"[ACCUMULATOR_DEBUG] Returning after handle_async_error")
                     return
                 # Send window EOF response or Window result response
                 # back to the client
                 else:
-                    logging.info(f"[ACCUMULATOR_DEBUG] Yielding message: {msg}")
                     yield msg
         except BaseException as e:
-            logging.info(f"[ACCUMULATOR_DEBUG] Caught exception in try block: {e}")
             await handle_async_error(context, e, ERR_UDF_EXCEPTION_STRING)
             return
         # Wait for the process_input_stream task to finish for a clean exit

pynumaflow/accumulator/servicer/task_manager.py

@@ -167,7 +167,7 @@ async def create_task(self, req):
 
             # Save the result of the reduce operation to the task list
             self.tasks[unified_key] = curr_task
-            
+
             # Increment expected EOF count since we created a new task
             async with self._eof_count_lock:
                 self._expected_eof_count += 1
@@ -277,7 +277,7 @@ async def process_input_stream(
             # Wait for all tasks to send their EOF responses before terminating the stream
             # This ensures proper ordering: all messages -> all EOF responses -> STREAM_EOF
             await self._stream_termination_event.wait()
-            
+
             # Now send STREAM_EOF to terminate the global result queue iterator
             await self.global_result_queue.put(STREAM_EOF)
         except BaseException as e:
@@ -309,17 +309,17 @@ async def write_to_global_queue(
             event_time_pb = timestamp_pb2.Timestamp()
             if msg.event_time is not None:
                 event_time_pb.FromDatetime(msg.event_time)
-            
+
             watermark_pb = timestamp_pb2.Timestamp()
             if msg.watermark is not None:
                 watermark_pb.FromDatetime(msg.watermark)
-            
+
             start_dt_pb = timestamp_pb2.Timestamp()
             start_dt_pb.FromDatetime(datetime.fromtimestamp(0))
-            
+
             end_dt_pb = timestamp_pb2.Timestamp()
             end_dt_pb.FromDatetime(wm)
-            
+
             res = accumulator_pb2.AccumulatorResponse(
                 payload=accumulator_pb2.Payload(
                     keys=msg.keys,
@@ -339,18 +339,18 @@ async def write_to_global_queue(
         # send EOF
         start_eof_pb = timestamp_pb2.Timestamp()
         start_eof_pb.FromDatetime(datetime.fromtimestamp(0))
-        
+
         end_eof_pb = timestamp_pb2.Timestamp()
         end_eof_pb.FromDatetime(wm)
-        
+
         res = accumulator_pb2.AccumulatorResponse(
             window=accumulator_pb2.KeyedWindow(
                 start=start_eof_pb, end=end_eof_pb, slot="slot-0", keys=task.keys
             ),
             EOF=True,
         )
         await output_queue.put(res)
-        
+
         # Increment received EOF count and check if all tasks are done
         async with self._eof_count_lock:
             self._received_eof_count += 1

tests/accumulator/test_async_accumulator.py

@@ -33,12 +33,14 @@ def request_generator(count, request, resetkey: bool = False):
             # Clear previous keys and add new ones
             del request.payload.keys[:]
             request.payload.keys.extend([f"key-{i}"])
-        
+
         # Set operation based on index - first is OPEN, rest are APPEND
         if i == 0:
             request.operation.event = accumulator_pb2.AccumulatorRequest.WindowOperation.Event.OPEN
         else:
-            request.operation.event = accumulator_pb2.AccumulatorRequest.WindowOperation.Event.APPEND
+            request.operation.event = (
+                accumulator_pb2.AccumulatorRequest.WindowOperation.Event.APPEND
+            )
         yield request
 
 
@@ -82,27 +84,19 @@ class ExampleClass(Accumulator):
     def __init__(self, counter):
         self.counter = counter
 
-    async def handler(
-        self, datums: AsyncIterable[Datum], output: NonBlockingIterator
-    ):
+    async def handler(self, datums: AsyncIterable[Datum], output: NonBlockingIterator):
         async for datum in datums:
             self.counter += 1
             msg = f"counter:{self.counter}"
-            await output.put(
-                Message(str.encode(msg), keys=datum.keys(), tags=[])
-            )
+            await output.put(Message(str.encode(msg), keys=datum.keys(), tags=[]))
 
 
-async def accumulator_handler_func(
-    datums: AsyncIterable[Datum], output: NonBlockingIterator
-):
+async def accumulator_handler_func(datums: AsyncIterable[Datum], output: NonBlockingIterator):
     counter = 0
     async for datum in datums:
         counter += 1
         msg = f"counter:{counter}"
-        await output.put(
-            Message(str.encode(msg), keys=datum.keys(), tags=[])
-        )
+        await output.put(Message(str.encode(msg), keys=datum.keys(), tags=[]))
 
 
 def NewAsyncAccumulator():
@@ -167,7 +161,7 @@ def test_accumulate(self) -> None:
         count = 0
         eof_count = 0
         for r in generator_response:
-            if hasattr(r, 'payload') and r.payload.value:
+            if hasattr(r, "payload") and r.payload.value:
                 count += 1
                 # Each datum should increment the counter
                 expected_msg = f"counter:{count}"
@@ -181,7 +175,7 @@ def test_accumulate(self) -> None:
             else:
                 self.assertEqual(r.EOF, True)
                 eof_count += 1
-                
+
         # We should have received 5 messages (one for each datum)
         self.assertEqual(5, count)
         self.assertEqual(1, eof_count)
@@ -209,7 +203,7 @@ def test_accumulate_with_multiple_keys(self) -> None:
                 # Track count per key
                 key = r.payload.keys[0] if r.payload.keys else "no_key"
                 key_counts[key] = key_counts.get(key, 0) + 1
-                
+
                 # Each key should have its own counter starting from 1
                 expected_msg = f"counter:{key_counts[key]}"
                 self.assertEqual(
@@ -220,7 +214,7 @@ def test_accumulate_with_multiple_keys(self) -> None:
             else:
                 eof_count += 1
                 self.assertEqual(r.EOF, True)
-                
+
         # We should have 10 messages (one for each key)
         self.assertEqual(10, count)
         self.assertEqual(10, eof_count)  # Each key/task sends its own EOF

tests/accumulator/test_async_accumulator_err.py

@@ -19,8 +19,6 @@
 from pynumaflow.shared.asynciter import NonBlockingIterator
 from tests.testing_utils import (
     mock_message,
-    mock_interval_window_start,
-    mock_interval_window_end,
     get_time_args,
     mock_terminate_on_stop,
 )
@@ -74,33 +72,25 @@ class ExampleErrorClass(Accumulator):
     def __init__(self, counter):
         self.counter = counter
 
-    async def handler(
-        self, datums: AsyncIterable[Datum], output: NonBlockingIterator
-    ):
+    async def handler(self, datums: AsyncIterable[Datum], output: NonBlockingIterator):
         async for datum in datums:
             self.counter += 1
             if self.counter == 2:
                 # Simulate an error on the second datum
                 raise RuntimeError("Simulated error in accumulator handler")
             msg = f"counter:{self.counter}"
-            await output.put(
-                Message(str.encode(msg), keys=datum.keys(), tags=[])
-            )
+            await output.put(Message(str.encode(msg), keys=datum.keys(), tags=[]))
 
 
-async def error_accumulator_handler_func(
-    datums: AsyncIterable[Datum], output: NonBlockingIterator
-):
+async def error_accumulator_handler_func(datums: AsyncIterable[Datum], output: NonBlockingIterator):
     counter = 0
     async for datum in datums:
         counter += 1
         if counter == 2:
             # Simulate an error on the second datum
             raise RuntimeError("Simulated error in accumulator function")
         msg = f"counter:{counter}"
-        await output.put(
-            Message(str.encode(msg), keys=datum.keys(), tags=[])
-        )
+        await output.put(Message(str.encode(msg), keys=datum.keys(), tags=[]))
 
 
 def NewAsyncAccumulatorError():
@@ -121,6 +111,7 @@ async def start_server(udfs):
     await server.start()
     await server.wait_for_termination()
 
+
 @patch("psutil.Process.kill", mock_terminate_on_stop)
 class TestAsyncAccumulatorError(unittest.TestCase):
     @classmethod
@@ -156,15 +147,15 @@ def test_accumulate_error(self) -> None:
         stub = self.__stub()
         request = start_request()
         generator_response = None
-        
+
         try:
             generator_response = stub.AccumulateFn(
                 request_iterator=request_generator(count=5, request=request)
             )
-            
+
             # Try to consume the generator
             counter = 0
-            logging.info(f"[TEST_DEBUG] About to iterate through generator_response")
+            logging.info("[TEST_DEBUG] About to iterate through generator_response")
             for response in generator_response:
                 counter += 1
                 logging.info(f"[TEST_DEBUG] Received response {counter}: {response}")
@@ -180,12 +171,12 @@ def test_accumulate_partial_success(self) -> None:
         """Test that the first datum is processed before error occurs"""
         stub = self.__stub()
         request = start_request()
-        
+
         try:
             generator_response = stub.AccumulateFn(
                 request_iterator=request_generator(count=3, request=request)
             )
-            
+
             # Try to consume the generator
             counter = 0
             for _ in generator_response: