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

* update

* remove py.typed for now

* update

* fix cimply and add types to serde producer/consumer

* admin

* address feedback

* add warning to stub and c files; admin typing more

* add accidentally removed md files

* fix merge conflicts in md files, add types to admin and serialization entrypoint init files

* finish admin init

* add types for AIO module

* linter fix

* address mypy complaints

* revert some accidental doc change

* fix some suggestions by copilot

* linter

* fix

* resolve conflict

* encryption clients

* fix

* revert incorrect merge conflict changes

* fix many things

* more fixes in non sr modules

* type encrypt_executor.py

* more typeignore removals

* update

* handle union types in schemas

* a bit more

* revert some bad changes during merge, address copilot comments

* minor

* support type hint substitution for unasync

* linter

* fix failign integration tests

* fix incorrect cache pattern

* fix type

* linter for comment

* pin google-re2<1.1.20251105 due to dropped python 3.9 support

Author: fangnx

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-rules.txt

@@ -2,6 +2,8 @@ azure-identity
 azure-keyvault-keys
 boto3>=1.35
 cel-python>=0.4.0
+# Pin google-re2 to last version with Python 3.9 wheels (see https://pypi.org/project/google-re2/1.1.20251105/#files)
+google-re2<1.1.20251105
 google-auth
 google-api-core
 google-cloud-kms

requirements/requirements-tests.txt

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

src/confluent_kafka/admin/__init__.py

@@ -649,11 +649,9 @@ def delete_topics(  # type: ignore[override]
         return futmap
 
     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: Optional[List[TopicPartition]]) -> None:
         self.topic_partitions = topic_partitions or []
 
 

src/confluent_kafka/admin/_resource.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
 from enum import Enum
 from .. import cimpl as _cimpl
 
@@ -28,7 +27,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 +43,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/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.,