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

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/_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

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,7 +28,7 @@ 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:
+    def __lt__(self, other: object) -> bool:
         if not isinstance(other, ResourceType):
             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:
+    def __lt__(self, other: object) -> bool:
         if not isinstance(other, ResourcePatternType):
             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

src/confluent_kafka/avro/cached_schema_registry_client.py

@@ -32,7 +32,7 @@
 
 # Python 2 considers int an instance of str
 try:
-    string_type = basestring  # noqa
+    string_type = basestring  # type: ignore[name-defined]  # noqa
 except NameError:
     string_type = str
 

src/confluent_kafka/avro/load.py

@@ -47,11 +47,11 @@ def _hash_func(self):
         from avro.errors import SchemaParseException
     except ImportError:
         # avro < 1.11.0
-        from avro.schema import SchemaParseException
+        from avro.schema import SchemaParseException  # type: ignore[attr-defined,no-redef]
 
-    schema.RecordSchema.__hash__ = _hash_func
-    schema.PrimitiveSchema.__hash__ = _hash_func
-    schema.UnionSchema.__hash__ = _hash_func
+    schema.RecordSchema.__hash__ = _hash_func  # type: ignore[method-assign]
+    schema.PrimitiveSchema.__hash__ = _hash_func  # type: ignore[method-assign]
+    schema.UnionSchema.__hash__ = _hash_func  # type: ignore[method-assign]
 
 except ImportError:
-    schema = None
+    schema = None # type: ignore[assignment]

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/_kafka_batch_executor.py

@@ -110,22 +110,22 @@ def _produce_batch_and_poll() -> int:
 
     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

@@ -72,58 +72,82 @@
 ]
 
 
-def topic_subject_name_strategy(ctx, record_name: Optional[str]) -> Optional[str]:
+def topic_subject_name_strategy(ctx: Optional[SerializationContext], record_name: Optional[str]) -> Optional[str]:
     """
     Constructs a subject name in the form of {topic}-key|value.
 
     Args:
         ctx (SerializationContext): Metadata pertaining to the serialization
-            operation.
+            operation. **Required** - will raise ValueError if None.
 
         record_name (Optional[str]): Record name.
 
+    Raises:
+        ValueError: If ctx is None.
+
     """
+    if ctx is None:
+        raise ValueError(
+            "SerializationContext is required for topic_subject_name_strategy. "
+            "Either provide a SerializationContext or use record_subject_name_strategy."
+        )
     return ctx.topic + "-" + ctx.field
 
 
-def topic_record_subject_name_strategy(ctx, record_name: Optional[str]) -> Optional[str]:
+def topic_record_subject_name_strategy(ctx: Optional[SerializationContext], record_name: Optional[str]) -> Optional[str]:
     """
     Constructs a subject name in the form of {topic}-{record_name}.
 
     Args:
         ctx (SerializationContext): Metadata pertaining to the serialization
-            operation.
+            operation. **Required** - will raise ValueError if None.
 
         record_name (Optional[str]): Record name.
 
+    Raises:
+        ValueError: If ctx is None.
+
     """
+    if ctx is None:
+        raise ValueError(
+            "SerializationContext is required for topic_record_subject_name_strategy. "
+            "Either provide a SerializationContext or use record_subject_name_strategy."
+        )
     return ctx.topic + "-" + record_name if record_name is not None else None
 
 
-def record_subject_name_strategy(ctx, record_name: Optional[str]) -> Optional[str]:
+def record_subject_name_strategy(ctx: Optional[SerializationContext], record_name: Optional[str]) -> Optional[str]:
     """
     Constructs a subject name in the form of {record_name}.
 
     Args:
         ctx (SerializationContext): Metadata pertaining to the serialization
-            operation.
+            operation. **Not used** by this strategy.
 
         record_name (Optional[str]): Record name.
 
+    Note:
+        This strategy does not require SerializationContext and can be used
+        when ctx is None.
+
     """
     return record_name if record_name is not None else None
 
 
-def reference_subject_name_strategy(ctx, schema_ref: SchemaReference) -> Optional[str]:
+def reference_subject_name_strategy(ctx: Optional[SerializationContext], schema_ref: SchemaReference) -> Optional[str]:
     """
     Constructs a subject reference name in the form of {reference name}.
 
     Args:
         ctx (SerializationContext): Metadata pertaining to the serialization
-            operation.
+            operation. **Not used** by this strategy.
 
         schema_ref (SchemaReference): SchemaReference instance.
 
+    Note:
+        This strategy does not require SerializationContext and can be used
+        when ctx is None.
+
     """
     return schema_ref.name if schema_ref is not None else None
 
@@ -205,7 +229,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