add types for AIO module

Author: fangnx

src/confluent_kafka/experimental/aio/_AIOConsumer.py

@@ -14,17 +14,18 @@
 
 import asyncio
 import concurrent.futures
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple
 
 import confluent_kafka
 
 from . import _common as _common
+from ..._types import ConfigDict
 
 
 class AIOConsumer:
     def __init__(
         self,
-        consumer_conf: Dict[str, Any],
+        consumer_conf: ConfigDict,
         max_workers: int = 2,
         executor: Optional[concurrent.futures.Executor] = None
     ) -> None:

src/confluent_kafka/experimental/aio/_common.py

@@ -18,6 +18,8 @@
 import concurrent.futures
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar
 
+from ..._types import ConfigDict
+
 T = TypeVar('T')
 
 
@@ -32,7 +34,7 @@ def __init__(
         self.logger = logger
 
     def log(self, *args: Any, **kwargs: Any) -> None:
-        self.loop.call_soon_threadsafe(lambda: self.logger.log(*args, **kwargs))
+        self.loop.call_soon_threadsafe(self.logger.log, *args, **kwargs)
 
 
 def wrap_callback(

src/confluent_kafka/experimental/aio/producer/_AIOProducer.py

@@ -15,14 +15,15 @@
 import asyncio
 import concurrent.futures
 import logging
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Optional
 
 import confluent_kafka
 
 from .. import _common as _common
 from ._producer_batch_processor import ProducerBatchManager
 from ._kafka_batch_executor import ProducerBatchExecutor
 from ._buffer_timeout_manager import BufferTimeoutManager
+from ..._types import ConfigDict
 
 
 logger = logging.getLogger(__name__)
@@ -36,7 +37,7 @@ class AIOProducer:
 
     def __init__(
         self,
-        producer_conf: Dict[str, Any],
+        producer_conf: ConfigDict,
         max_workers: int = 4,
         executor: Optional[concurrent.futures.Executor] = None,
         batch_size: int = 1000,
@@ -224,7 +225,7 @@ async def flush(self, *args: Any, **kwargs: Any) -> Any:
             # Update buffer activity since we just flushed
             self._buffer_timeout_manager.mark_activity()
 
-        # Then flush the underlying producer and wait for delivery confirmation
+        # Then flush underlying producer and wait for delivery confirmation
         return await self._call(self._producer.flush, *args, **kwargs)
 
     async def purge(self, *args: Any, **kwargs: Any) -> Any:

src/confluent_kafka/experimental/aio/producer/_kafka_batch_executor.py

@@ -19,8 +19,6 @@
 
 import confluent_kafka
 
-from .. import _common
-
 logger = logging.getLogger(__name__)
 
 
@@ -65,10 +63,10 @@ async def execute_batch(
         Args:
             topic: Target topic for the batch
             batch_messages: List of prepared messages with callbacks assigned
-            partition: Target partition for the batch (-1 = RD_KAFKA_PARTITION_UA)
+            partition: Target partition (-1 = RD_KAFKA_PARTITION_UA)
 
         Returns:
-            Result from producer.poll() indicating number of delivery reports processed
+            Result from producer.poll() indicating # of delivery reports processed
 
         Raises:
             Exception: Any exception from the batch operation is propagated
@@ -77,9 +75,9 @@ def _produce_batch_and_poll() -> int:
             """Helper function to run in thread pool
 
             This function encapsulates all the blocking Kafka operations:
-            - Call produce_batch with specific partition and individual message callbacks
+            - Call produce_batch with specific partition & individual callbacks
             - Handle partial batch failures for messages that fail immediately
-            - Poll for delivery reports to trigger callbacks for successful messages
+            - Poll for delivery reports to trigger callbacks for successful msgs
             """
             # Call produce_batch with specific partition and individual callbacks
             # Convert tuple to list since produce_batch expects a list
@@ -90,7 +88,8 @@ def _produce_batch_and_poll() -> int:
             )
 
             # Use the provided partition for the entire batch
-            # This enables proper partition control while working around librdkafka limitations
+            # This enables proper partition control while working around
+            # librdkafka limitations
             self._producer.produce_batch(topic, messages_list, partition=partition)
 
             # Handle partial batch failures: Check for messages that failed
@@ -99,7 +98,7 @@ def _produce_batch_and_poll() -> int:
             # so we need to manually invoke their callbacks
             self._handle_partial_failures(messages_list)
 
-            # Immediately poll to process delivery callbacks for successful messages
+            # Immediately poll to process delivery callbacks for successful msgs
             poll_result = self._producer.poll(0)
 
             return poll_result
@@ -108,24 +107,10 @@ def _produce_batch_and_poll() -> int:
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(self._executor, _produce_batch_and_poll)
 
-    async def flush_librdkafka_queue(self, timeout=-1):
-        """Flush the librdkafka queue and wait for all messages to be delivered
-
-        This method awaits until all outstanding produce requests are completed
-        or the timeout is reached, unless the timeout is set to 0 (non-blocking).
-
-        Args:
-            timeout: Maximum time to wait in seconds:
-                    - -1 = wait indefinitely (default)
-                    - 0 = non-blocking, return immediately
-                    - >0 = wait up to timeout seconds
-
-        Returns:
-            Number of messages still in queue after flush attempt
-        """
-        return await _common.async_call(self._executor, self._producer.flush, timeout)
-
-    def _handle_partial_failures(self, batch_messages: List[Dict[str, Any]]) -> None:
+    def _handle_partial_failures(
+        self,
+        batch_messages: List[Dict[str, Any]]
+    ) -> None:
         """Handle messages that failed during produce_batch
 
         When produce_batch encounters messages that fail immediately (e.g.,
@@ -135,7 +120,7 @@ def _handle_partial_failures(self, batch_messages: List[Dict[str, Any]]) -> None
         manually invoke the simple future-resolving callbacks.
 
         Args:
-            batch_messages: List of message dictionaries that were passed to produce_batch
+            batch_messages: List of message dicts passed to produce_batch
         """
         for msg_dict in batch_messages:
             if '_error' in msg_dict:
@@ -146,7 +131,7 @@ def _handle_partial_failures(self, batch_messages: List[Dict[str, Any]]) -> None
                     # Extract the error from the message dict (set by Producer.c)
                     error = msg_dict['_error']
                     # Manually invoke the callback with the error
-                    # Note: msg is None since the message failed before being queued
+                    # Note: msg is None since message failed before being queued
                     try:
                         callback(error, None)
                     except Exception: