Ensure stubtests for kafka-python passes

Author: RasmusNygren

stubs/kafka-python/@tests/stubtest_allowlist.txt

@@ -0,0 +1,21 @@
+# Command-line entry points are not a useful typed API surface.
+kafka.__main__
+kafka.admin.__main__
+kafka.consumer.__main__
+kafka.producer.__main__
+
+# Concrete subclasses define these abstract properties as class attributes.
+kafka.protocol.api.Request.API_KEY
+kafka.protocol.api.Request.API_VERSION
+kafka.protocol.api.Request.RESPONSE_TYPE
+kafka.protocol.api.Request.SCHEMA
+kafka.protocol.api.Response.API_KEY
+kafka.protocol.api.Response.API_VERSION
+kafka.protocol.api.Response.SCHEMA
+
+# Vendored compatibility modules are implementation details.
+kafka.vendor
+kafka.vendor.enum34
+kafka.vendor.selectors34
+kafka.vendor.six
+kafka.vendor.socketpair

stubs/kafka-python/METADATA.toml

@@ -1,2 +1,5 @@
 version = "2.3.*"
 upstream-repository = "https://github.com/dpkp/kafka-python"
+
+[tool.stubtest]
+stubtest-dependencies = ["pyperf"]

stubs/kafka-python/kafka/admin/client.pyi

@@ -2,7 +2,7 @@ import selectors
 import ssl
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterable, Mapping, Sequence
-from typing import Literal, TypeAlias, TypedDict
+from typing import Literal, TypeAlias, TypedDict, type_check_only
 from typing_extensions import Unpack
 
 from kafka.admin.acl_resource import ACL, ACLFilter
@@ -20,6 +20,7 @@ _SaslMechanism: TypeAlias = Literal["PLAIN", "GSSAPI", "OAUTHBEARER", "SCRAM-SHA
 _SecurityProtocol: TypeAlias = Literal["PLAINTEXT", "SSL", "SASL_PLAINTEXT", "SASL_SSL"]
 _SocketOption: TypeAlias = tuple[int, int, int]
 
+@type_check_only
 class _KafkaAdminClientConfig(TypedDict, total=False):
     bootstrap_servers: _BootstrapServers
     client_id: str
@@ -59,6 +60,7 @@ class _KafkaAdminClientConfig(TypedDict, total=False):
     metrics_sample_window_ms: int
     kafka_client: _KafkaClientFactory
 
+@type_check_only
 class _CreateAclsResult(TypedDict):
     succeeded: list[ACL]
     failed: list[tuple[ACL, KafkaError]]

stubs/kafka-python/kafka/client_async.pyi

@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from collections import OrderedDict as OrderedDict
 
 log: Incomplete
 socketpair: Incomplete
@@ -33,8 +34,6 @@ class KafkaClient:
     def await_ready(self, node_id, timeout_ms: int = 30000): ...
     def send_and_receive(self, node_id, request): ...
 
-OrderedDict = dict
-
 class IdleConnectionManager:
     connections_max_idle: Incomplete
     next_idle_close_check_time: Incomplete

stubs/kafka-python/kafka/consumer/fetcher.pyi

@@ -1,5 +1,5 @@
 from _typeshed import Incomplete
-from typing import NamedTuple
+from typing import ClassVar, NamedTuple
 
 import kafka.errors as Errors
 
@@ -99,6 +99,8 @@ class FetchMetadata:
     THROTTLED_SESSION_ID: int
     INITIAL_EPOCH: int
     FINAL_EPOCH: int
+    INITIAL: ClassVar[FetchMetadata]
+    LEGACY: ClassVar[FetchMetadata]
     session_id: Incomplete
     epoch: Incomplete
     def __init__(self, session_id, epoch) -> None: ...

stubs/kafka-python/kafka/consumer/group.pyi

@@ -2,7 +2,7 @@ import selectors
 import ssl
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
-from typing import Literal, TypeAlias, TypedDict, overload
+from typing import Literal, TypeAlias, TypedDict, overload, type_check_only
 from typing_extensions import Self, Unpack
 
 from kafka.consumer.fetcher import ConsumerRecord
@@ -20,6 +20,7 @@ _SaslMechanism: TypeAlias = Literal["PLAIN", "GSSAPI", "OAUTHBEARER", "SCRAM-SHA
 _SecurityProtocol: TypeAlias = Literal["PLAINTEXT", "SSL", "SASL_PLAINTEXT", "SASL_SSL"]
 _SocketOption: TypeAlias = tuple[int, int, int]
 
+@type_check_only
 class _KafkaConsumerConfig(TypedDict, total=False):
     bootstrap_servers: _BootstrapServers
     client_id: str

stubs/kafka-python/kafka/coordinator/base.pyi

@@ -1,6 +1,7 @@
 import abc
 import threading
 from _typeshed import Incomplete
+from typing import ClassVar
 
 from kafka import errors as Errors
 
@@ -13,6 +14,7 @@ class MemberState:
     STABLE: str
 
 class Generation:
+    NO_GENERATION: ClassVar[Generation]
     generation_id: Incomplete
     member_id: Incomplete
     protocol: Incomplete

stubs/kafka-python/kafka/producer/kafka.pyi

@@ -2,7 +2,7 @@ import selectors
 import ssl
 from _typeshed import Incomplete
 from collections.abc import Callable, Mapping, Sequence
-from typing import Literal, TypeAlias, TypedDict
+from typing import Literal, TypeAlias, TypedDict, type_check_only
 from typing_extensions import Unpack
 
 from kafka.producer.future import FutureRecordMetadata
@@ -18,6 +18,7 @@ _SaslMechanism: TypeAlias = Literal["PLAIN", "GSSAPI", "OAUTHBEARER", "SCRAM-SHA
 _SecurityProtocol: TypeAlias = Literal["PLAINTEXT", "SSL", "SASL_PLAINTEXT", "SASL_SSL"]
 _SocketOption: TypeAlias = tuple[int, int, int]
 
+@type_check_only
 class _KafkaProducerConfig(TypedDict, total=False):
     bootstrap_servers: _BootstrapServers
     client_id: str | None

stubs/kafka-python/kafka/producer/sender.pyi

@@ -5,14 +5,14 @@ from typing import NamedTuple
 log: Incomplete
 
 class PartitionResponse(NamedTuple):
-    error: Incomplete
-    base_offset: Incomplete
-    last_offset: Incomplete
-    log_append_time: Incomplete
-    log_start_offset: Incomplete
-    record_errors: Incomplete
-    error_message: Incomplete
-    current_leader: Incomplete
+    error: Incomplete = ...
+    base_offset: Incomplete = ...
+    last_offset: Incomplete = ...
+    log_append_time: Incomplete = ...
+    log_start_offset: Incomplete = ...
+    record_errors: Incomplete = ...
+    error_message: Incomplete = ...
+    current_leader: Incomplete = ...
 
 class Sender(threading.Thread):
     DEFAULT_CONFIG: Incomplete

stubs/kafka-python/kafka/protocol/group.pyi

@@ -8,9 +8,9 @@ DEFAULT_GENERATION_ID: int
 UNKNOWN_MEMBER_ID: str
 
 class GroupMember(NamedTuple):
-    member_id: Incomplete
-    group_instance_id: Incomplete
-    metadata_bytes: Incomplete
+    member_id: Incomplete = ...
+    group_instance_id: Incomplete = ...
+    metadata_bytes: Incomplete = ...
 
 class JoinGroupResponse_v0(Response):
     API_KEY: int