Set up type hinting: add correct types and apply type fixes in schema-registry module

pyproject.toml

@@ -27,6 +27,21 @@ Homepage = "https://github.com/confluentinc/confluent-kafka-python"
 [tool.mypy]
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = [
+    "confluent_kafka.schema_registry.avro",
+    "confluent_kafka.schema_registry.json_schema",
+    "confluent_kafka.schema_registry.protobuf",
+]
+disable_error_code = ["assignment", "no-redef"]
+
+[[tool.mypy.overrides]]
+module = [
+    "confluent_kafka.schema_registry.confluent.meta_pb2",
+    "confluent_kafka.schema_registry.confluent.types.decimal_pb2",
+]
+ignore_errors = true
+
 [tool.setuptools]
 include-package-data = false
 

requirements/requirements-tests.txt

@@ -1,6 +1,8 @@
 # core test requirements
 urllib3<3
 flake8
+mypy
+types-cachetools
 orjson
 pytest
 pytest-timeout

src/confluent_kafka/_model/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import List, Optional, Any
 from enum import Enum
 from .. import cimpl
 from ..cimpl import TopicPartition
@@ -91,8 +91,8 @@ class ConsumerGroupState(Enum):
     #: Consumer Group is empty.
     EMPTY = cimpl.CONSUMER_GROUP_STATE_EMPTY
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ConsumerGroupState):
+    def __lt__(self, other) -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -111,8 +111,8 @@ class ConsumerGroupType(Enum):
     #: Classic Type
     CLASSIC = cimpl.CONSUMER_GROUP_TYPE_CLASSIC
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ConsumerGroupType):
+    def __lt__(self, other) -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -167,8 +167,8 @@ class IsolationLevel(Enum):
     READ_UNCOMMITTED = cimpl.ISOLATION_LEVEL_READ_UNCOMMITTED  #: Receive all the offsets.
     READ_COMMITTED = cimpl.ISOLATION_LEVEL_READ_COMMITTED  #: Skip offsets belonging to an aborted transaction.
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, IsolationLevel):
+    def __lt__(self, other) -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -186,7 +186,7 @@ class ElectionType(Enum):
     #: Unclean election
     UNCLEAN = cimpl.ELECTION_TYPE_UNCLEAN
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ElectionType):
+    def __lt__(self, other) -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value

src/confluent_kafka/admin/__init__.py

@@ -653,7 +653,6 @@ def list_topics(self, *args: Any, **kwargs: Any) -> ClusterMetadata:
         return super(AdminClient, self).list_topics(*args, **kwargs)
 
     def list_groups(self, *args: Any, **kwargs: Any) -> List[GroupMetadata]:
-
         return super(AdminClient, self).list_groups(*args, **kwargs)
 
     def create_partitions(  # type: ignore[override]

src/confluent_kafka/admin/_acl.py

@@ -44,8 +44,8 @@ class AclOperation(Enum):
     ALTER_CONFIGS = _cimpl.ACL_OPERATION_ALTER_CONFIGS  #: ALTER_CONFIGS operation
     IDEMPOTENT_WRITE = _cimpl.ACL_OPERATION_IDEMPOTENT_WRITE  #: IDEMPOTENT_WRITE operation
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, AclOperation):
+    def __lt__(self, other: 'AclOperation') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -59,8 +59,8 @@ class AclPermissionType(Enum):
     DENY = _cimpl.ACL_PERMISSION_TYPE_DENY  #: Disallows access
     ALLOW = _cimpl.ACL_PERMISSION_TYPE_ALLOW  #: Grants access
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, AclPermissionType):
+    def __lt__(self, other: 'AclPermissionType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -161,8 +161,8 @@ def _to_tuple(self) -> Tuple[ResourceType, str, ResourcePatternType, str, str, A
     def __hash__(self) -> int:
         return hash(self._to_tuple())
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, AclBinding):
+    def __lt__(self, other: 'AclBinding') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self._to_tuple() < other._to_tuple()
 

src/confluent_kafka/admin/_config.py

@@ -185,9 +185,7 @@ def __repr__(self) -> str:
     def __hash__(self) -> int:
         return hash((self.restype, self.name))
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ConfigResource):
-            return NotImplemented
+    def __lt__(self, other: 'ConfigResource') -> bool:
         if self.restype < other.restype:
             return True
         return self.name.__lt__(other.name)

src/confluent_kafka/admin/_group.py

@@ -80,7 +80,7 @@ class MemberAssignment:
         The topic partitions assigned to a group member.
     """
 
-    def __init__(self, topic_partitions: Optional[List[TopicPartition]] = None) -> None:
+    def __init__(self, topic_partitions: List[TopicPartition] = []) -> None:
         self.topic_partitions = topic_partitions or []
 
 

src/confluent_kafka/admin/_listoffsets.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Optional
+from typing import Dict, Any, Optional
 from abc import ABC, abstractmethod
 
 from .. import cimpl
@@ -68,7 +68,7 @@ def __new__(cls, index: int):
         else:
             return cls.for_timestamp(index)
 
-    def __lt__(self, other: object) -> bool:
+    def __lt__(self, other) -> Any:
         if not isinstance(other, OffsetSpec):
             return NotImplemented
         return self._value < other._value

src/confluent_kafka/admin/_metadata.py

@@ -79,7 +79,7 @@ class TopicMetadata(object):
     # on other classes which raises a warning/error.
 
     def __init__(self) -> None:
-        self.topic: Optional[str] = None
+        self.topic = None
         """Topic name"""
         self.partitions: Dict[int, 'PartitionMetadata'] = {}
         """Map of partitions indexed by partition id. Value is a PartitionMetadata object."""

src/confluent_kafka/admin/_resource.py

@@ -28,8 +28,8 @@ class ResourceType(Enum):
     BROKER = _cimpl.RESOURCE_BROKER  #: Broker resource. Resource name is broker id.
     TRANSACTIONAL_ID = _cimpl.RESOURCE_TRANSACTIONAL_ID  #: Transactional ID resource.
 
-    def __lt__(self, other: object) -> Any:
-        if not isinstance(other, ResourceType):
+    def __lt__(self, other: 'ResourceType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 
@@ -44,7 +44,7 @@ class ResourcePatternType(Enum):
     LITERAL = _cimpl.RESOURCE_PATTERN_LITERAL  #: Literal: A literal resource name
     PREFIXED = _cimpl.RESOURCE_PATTERN_PREFIXED  #: Prefixed: A prefixed resource name
 
-    def __lt__(self, other: object) -> Any:
-        if not isinstance(other, ResourcePatternType):
+    def __lt__(self, other: 'ResourcePatternType') -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value

src/confluent_kafka/admin/_scram.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import List, Optional, Any
 from enum import Enum
 
 from .. import cimpl
@@ -26,8 +26,8 @@ class ScramMechanism(Enum):
     SCRAM_SHA_256 = cimpl.SCRAM_MECHANISM_SHA_256  #: SCRAM-SHA-256 mechanism
     SCRAM_SHA_512 = cimpl.SCRAM_MECHANISM_SHA_512  #: SCRAM-SHA-512 mechanism
 
-    def __lt__(self, other: object) -> bool:
-        if not isinstance(other, ScramMechanism):
+    def __lt__(self, other) -> Any:
+        if self.__class__ != other.__class__:
             return NotImplemented
         return self.value < other.value
 

src/confluent_kafka/cimpl.pyi

@@ -34,14 +34,13 @@ TODO: Consider migrating to Cython in the future to eliminate this dual
 maintenance burden and get type hints directly from the implementation.
 """
 
-from typing import Any, Optional, Callable, List, Tuple, Dict, Union, overload, TYPE_CHECKING
+from typing import Any, Optional, Callable, List, Tuple, Dict, Union, overload
 from typing_extensions import Self, Literal
 import builtins
 
-from ._types import HeadersType
+from confluent_kafka.admin._metadata import ClusterMetadata, GroupMetadata
 
-if TYPE_CHECKING:
-    from confluent_kafka.admin._metadata import ClusterMetadata, GroupMetadata
+from ._types import HeadersType
 
 # Callback types with proper class references (defined locally to avoid circular imports)
 DeliveryCallback = Callable[[Optional['KafkaError'], 'Message'], None]

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

@@ -224,7 +224,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__)
 
 
@@ -108,24 +106,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.,

src/confluent_kafka/schema_registry/__init__.py

@@ -205,7 +205,7 @@ def dual_schema_id_deserializer(payload: bytes, ctx: Optional[SerializationConte
 
     # Parse schema ID from determined source and return appropriate payload
     if header_value is not None:
-        schema_id.from_bytes(io.BytesIO(header_value))
+        schema_id.from_bytes(io.BytesIO(header_value))  # type: ignore[arg-type]
         return io.BytesIO(payload)  # Return full payload when schema ID is in header
     else:
         return schema_id.from_bytes(io.BytesIO(payload))  # Parse from payload, return remainder