kafka-python 3.0.0__py3-none-any.whl

kafka/protocol/generate_stubs.py

@@ -0,0 +1,468 @@
+"""Generate .pyi type stub files for dynamically created protocol classes.
+
+Usage:
+    python -m kafka.protocol.generate_stubs           # generate stubs
+    python -m kafka.protocol.generate_stubs --check    # exit 1 if stubs are out of date
+    python -m kafka.protocol.generate_stubs --dry-run  # print to stdout
+"""
+import importlib
+import inspect
+import os
+import sys
+
+from .api_data import ApiData
+from .api_message import ApiMessage
+from .data_container import DataContainer
+from .schemas.fields.simple import SimpleField
+from .schemas.fields.array import ArrayField
+from .schemas.fields.struct import StructField
+from .schemas.fields.struct_array import StructArrayField
+from kafka.util import classproperty
+
+
+# JSON schema type string -> Python type annotation
+SIMPLE_TYPE_MAP = {
+    'int8': 'int',
+    'int16': 'int',
+    'uint16': 'int',
+    'int32': 'int',
+    'int64': 'int',
+    'float64': 'float',
+    'bool': 'bool',
+    'string': 'str',
+    'bytes': 'bytes',
+    'records': 'bytes',
+    'uuid': 'uuid.UUID',
+    'bitfield': 'set[int]',
+}
+
+# Modules containing protocol classes, relative to kafka.protocol
+MESSAGE_MODULES = [
+    'kafka.protocol.admin.acl',
+    'kafka.protocol.admin.client_quotas',
+    'kafka.protocol.admin.cluster',
+    'kafka.protocol.admin.configs',
+    'kafka.protocol.admin.groups',
+    'kafka.protocol.admin.topics',
+    'kafka.protocol.admin.transactions',
+    'kafka.protocol.admin.users',
+    'kafka.protocol.consumer.fetch',
+    'kafka.protocol.consumer.group',
+    'kafka.protocol.consumer.metadata',
+    'kafka.protocol.consumer.offsets',
+    'kafka.protocol.metadata.api_versions',
+    'kafka.protocol.metadata.find_coordinator',
+    'kafka.protocol.metadata.metadata',
+    'kafka.protocol.producer.produce',
+    'kafka.protocol.producer.transaction',
+    'kafka.protocol.sasl',
+]
+
+
+def _all_fields_recursive(cls):
+    """Yield all fields from cls and its nested structs."""
+    if cls._struct is None:
+        return
+    for field in cls._struct._fields:
+        yield field
+        if field.is_struct() or field.is_struct_array():
+            dc = getattr(cls, field.type_str, None)
+            if dc is not None:
+                yield from _all_fields_recursive(dc)
+
+
+def resolve_type(field):
+    """Convert a field object to a Python type annotation string."""
+    if isinstance(field, StructArrayField):
+        base = 'list[%s]' % field.type_str
+    elif isinstance(field, ArrayField):
+        # Simple array (e.g., []int32, []string)
+        inner = SIMPLE_TYPE_MAP.get(field.array_of._type_str, 'Any')
+        base = 'list[%s]' % inner
+    elif isinstance(field, StructField):
+        base = field.type_str
+    elif isinstance(field, SimpleField):
+        base = SIMPLE_TYPE_MAP.get(field._type_str, 'Any')
+        # bytes fields accept ApiData objects (encoded automatically via .encode())
+        if field._type_str in ('bytes', 'records'):
+            base = 'bytes | ApiData'
+    else:
+        base = 'Any'
+
+    if field._nullable_versions is not None:
+        return '%s | None' % base
+    return base
+
+
+def _collect_nested_classes(cls):
+    """Yield (field, data_class) pairs for nested struct types on cls."""
+    if cls._struct is None:
+        return
+    seen = set()
+    for field in cls._struct._fields:
+        if field.is_struct() or field.is_struct_array():
+            dc = getattr(cls, field.type_str, None)
+            if dc is not None and dc not in seen:
+                seen.add(dc)
+                yield field, dc
+
+
+def emit_class(cls, indent=0, base_name=None):
+    """Generate stub lines for a single protocol class."""
+    pad = ' ' * indent
+    lines = []
+
+    # Determine base class name for the stub
+    if base_name is None:
+        if issubclass(cls, ApiMessage):
+            base_name = 'ApiMessage'
+        elif issubclass(cls, ApiData):
+            base_name = 'ApiData'
+        else:
+            base_name = 'DataContainer'
+
+    lines.append('%sclass %s(%s):' % (pad, cls.__name__, base_name))
+
+    inner_pad = pad + '    '
+
+    # Nested inner classes first
+    for field, dc in _collect_nested_classes(cls):
+        lines.extend(emit_class(dc, indent=indent + 4, base_name='DataContainer'))
+        lines.append('')
+
+    # Field annotations
+    if cls._struct is not None:
+        for field in cls._struct._fields:
+            type_ann = resolve_type(field)
+            lines.append('%s%s: %s' % (inner_pad, field.name, type_ann))
+
+    # __init__ signature
+    init_params = ['self', '*args: Any']
+    if cls._struct is not None:
+        for field in cls._struct._fields:
+            type_ann = resolve_type(field)
+            init_params.append('%s: %s = ...' % (field.name, type_ann))
+    init_params.append('version: int | None = None')
+    init_params.append('**kwargs: Any')
+
+    if len(init_params) <= 4:
+        lines.append('%sdef __init__(%s) -> None: ...' % (inner_pad, ', '.join(init_params)))
+    else:
+        lines.append('%sdef __init__(' % inner_pad)
+        for i, param in enumerate(init_params):
+            comma = ',' if i < len(init_params) - 1 else ','
+            lines.append('%s    %s%s' % (inner_pad, param, comma))
+        lines.append('%s) -> None: ...' % inner_pad)
+
+    # Common DataContainer methods/properties
+    lines.append('%s@property' % inner_pad)
+    lines.append('%sdef version(self) -> int | None: ...' % inner_pad)
+    lines.append('%sdef to_dict(self, meta: bool = False, json: bool = True) -> dict: ...' % inner_pad)
+
+    if base_name == 'DataContainer':
+        pass  # encode/decode inherited from DataContainer
+    elif base_name == 'ApiData':
+        _emit_api_data_methods(lines, inner_pad, cls)
+    elif base_name == 'ApiMessage':
+        _emit_api_message_methods(lines, inner_pad, cls)
+
+    # Detect custom members defined on the class itself (not inherited)
+    _emit_custom_members(lines, inner_pad, cls, base_name)
+
+    return lines
+
+
+def _emit_api_data_methods(lines, pad, cls):
+    """Emit method stubs for ApiData subclasses."""
+    lines.append('%sname: str' % pad)
+    lines.append('%stype: str' % pad)
+    lines.append('%svalid_versions: tuple[int, int]' % pad)
+    lines.append('%smin_version: int' % pad)
+    lines.append('%smax_version: int' % pad)
+
+
+def _emit_api_message_methods(lines, pad, cls):
+    """Emit method stubs for ApiMessage subclasses."""
+    lines.append('%sname: str' % pad)
+    lines.append('%stype: str' % pad)
+    lines.append('%sAPI_KEY: int' % pad)
+    lines.append('%sAPI_VERSION: int' % pad)
+    lines.append('%svalid_versions: tuple[int, int]' % pad)
+    lines.append('%smin_version: int' % pad)
+    lines.append('%smax_version: int' % pad)
+    lines.append('%s@property' % pad)
+    lines.append('%sdef header(self) -> Any: ...' % pad)
+    lines.append('%s@classmethod' % pad)
+    lines.append('%sdef is_request(cls) -> bool: ...' % pad)
+    lines.append('%sdef expect_response(self) -> bool: ...' % pad)
+    lines.append('%sdef with_header(self, correlation_id: int = 0, client_id: str = "kafka-python") -> None: ...' % pad)
+
+
+def _format_annotation(ann):
+    """Format a type annotation as a clean string, using short class names."""
+    if isinstance(ann, str):
+        return ann
+    if isinstance(ann, type):
+        return ann.__name__
+    # For generic aliases like list[TopicPartition], use repr and clean up module paths
+    s = str(ann)
+    # Remove module prefixes (e.g., kafka.structs.TopicPartition -> TopicPartition)
+    import re
+    s = re.sub(r'[a-zA-Z_][a-zA-Z0-9_.]*\.([A-Z][a-zA-Z0-9_]*)', r'\1', s)
+    return s
+
+
+def _get_return_annotation(obj):
+    """Extract the return type annotation string from a callable, property, or classproperty."""
+    func = None
+    if isinstance(obj, classproperty):
+        func = obj.f
+    elif isinstance(obj, property):
+        func = obj.fget
+    elif isinstance(obj, classmethod):
+        func = obj.__func__
+    elif isinstance(obj, staticmethod):
+        func = obj.__func__
+    elif callable(obj):
+        func = obj
+
+    if func is not None:
+        hints = getattr(func, '__annotations__', {})
+        ret = hints.get('return')
+        if ret is not None:
+            return _format_annotation(ret)
+    return None
+
+
+def _get_param_annotations(obj):
+    """Extract parameter annotations from a callable, returning (params_str, has_annotations)."""
+    func = None
+    if isinstance(obj, classmethod):
+        func = obj.__func__
+    elif isinstance(obj, staticmethod):
+        func = obj.__func__
+    elif callable(obj):
+        func = obj
+
+    if func is None:
+        return None
+
+    try:
+        sig = inspect.signature(func)
+    except (ValueError, TypeError):
+        return None
+
+    parts = []
+    for pname, param in sig.parameters.items():
+        if pname == 'self' or pname == 'cls':
+            parts.append(pname)
+            continue
+        ann = param.annotation
+        if ann is inspect.Parameter.empty:
+            ann_str = 'Any'
+        else:
+            ann_str = _format_annotation(ann)
+        if param.default is not inspect.Parameter.empty:
+            parts.append('%s: %s = ...' % (pname, ann_str))
+        elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+            parts.append('*%s: %s' % (pname, ann_str))
+        elif param.kind == inspect.Parameter.VAR_KEYWORD:
+            parts.append('**%s: %s' % (pname, ann_str))
+        else:
+            parts.append('%s: %s' % (pname, ann_str))
+    return ', '.join(parts)
+
+
+def _emit_custom_members(lines, pad, cls, base_name):
+    """Emit stubs for custom members defined directly on the class.
+
+    Reads type annotations from the source when available.
+    """
+    skip = {
+        '_json', '_struct', '_class_version', '_VERSIONS', '_flexible_versions',
+        '_valid_versions', '__module__', '__qualname__', '__doc__', '__dict__',
+        '__weakref__', '__license__', 'json_patch', '__init_subclass__',
+    }
+    if base_name == 'ApiMessage':
+        base = ApiMessage
+    elif base_name == 'ApiData':
+        base = ApiData
+    else:
+        base = DataContainer
+
+    for name, obj in cls.__dict__.items():
+        if name.startswith('_') or name in skip:
+            continue
+        # Skip encode/decode overrides to avoid LSP violations in stubs
+        if name in ('encode', 'decode', 'encode_into') and hasattr(base, name):
+            continue
+        # Skip slot descriptors and nested classes (already handled)
+        if isinstance(obj, type):
+            continue
+        if type(obj).__name__ == 'member_descriptor':
+            continue
+
+        ret = _get_return_annotation(obj) or 'Any'
+
+        if isinstance(obj, classproperty):
+            lines.append('%s@property' % pad)
+            lines.append('%sdef %s(self) -> %s: ...' % (pad, name, ret))
+        elif isinstance(obj, property):
+            lines.append('%s@property' % pad)
+            lines.append('%sdef %s(self) -> %s: ...' % (pad, name, ret))
+        elif isinstance(obj, classmethod):
+            params = _get_param_annotations(obj) or 'cls, *args: Any, **kwargs: Any'
+            lines.append('%s@classmethod' % pad)
+            lines.append('%sdef %s(%s) -> %s: ...' % (pad, name, params, ret))
+        elif isinstance(obj, staticmethod):
+            params = _get_param_annotations(obj) or '*args: Any, **kwargs: Any'
+            lines.append('%s@staticmethod' % pad)
+            lines.append('%sdef %s(%s) -> %s: ...' % (pad, name, params, ret))
+        elif callable(obj) and not isinstance(obj, type):
+            params = _get_param_annotations(obj) or 'self, *args: Any, **kwargs: Any'
+            lines.append('%sdef %s(%s) -> %s: ...' % (pad, name, params, ret))
+
+
+def discover_modules():
+    """Import all message modules and collect their exports.
+
+    Returns dict mapping module file path to list of (name, obj) pairs.
+    """
+    result = {}
+    for mod_name in MESSAGE_MODULES:
+        mod = importlib.import_module(mod_name)
+        mod_file = inspect.getfile(mod)
+        all_names = getattr(mod, '__all__', [])
+        exports = []
+        for name in all_names:
+            obj = getattr(mod, name)
+            exports.append((name, obj))
+        result[mod_file] = (mod, exports)
+    return result
+
+
+def generate_module(mod, exports):
+    """Generate the complete .pyi file content for a module."""
+    lines = []
+    lines.append('# Generated by generate_stubs.py (Python %d.%d)' % sys.version_info[:2])
+    lines.append('import uuid')
+    lines.append('from typing import Any, Self')
+    lines.append('')
+
+    # Determine needed imports based on base classes
+    from enum import IntEnum as _IntEnum
+    needs_api_message = False
+    needs_api_data = False
+    needs_data_container = False
+    needs_int_enum = False
+    for name, obj in exports:
+        if isinstance(obj, type):
+            if issubclass(obj, ApiMessage):
+                needs_api_message = True
+            elif issubclass(obj, ApiData):
+                needs_api_data = True
+            if issubclass(obj, _IntEnum):
+                needs_int_enum = True
+
+    # Check fields for nested classes (DataContainer) and bytes fields (ApiData)
+    for name, obj in exports:
+        if isinstance(obj, type) and hasattr(obj, '_struct') and obj._struct is not None:
+            for field in _all_fields_recursive(obj):
+                if field.is_struct() or field.is_struct_array():
+                    needs_data_container = True
+                if isinstance(field, SimpleField) and field._type_str in ('bytes', 'records'):
+                    needs_api_data = True
+
+    if needs_int_enum:
+        lines.append('from enum import IntEnum')
+    if needs_api_message:
+        lines.append('from kafka.protocol.api_message import ApiMessage')
+    if needs_api_data:
+        lines.append('from kafka.protocol.api_data import ApiData')
+    if needs_data_container:
+        lines.append('from kafka.protocol.data_container import DataContainer')
+    if needs_api_message or needs_api_data or needs_data_container or needs_int_enum:
+        lines.append('')
+
+    # Export __all__ so __init__.py can import it
+    all_names = [name for name, obj in exports]
+    lines.append('__all__ = [%s]' % ', '.join("'%s'" % n for n in all_names))
+    lines.append('')
+
+    for name, obj in exports:
+        if isinstance(obj, type) and issubclass(obj, DataContainer) and obj._struct is not None:
+            lines.extend(emit_class(obj))
+            lines.append('')
+        elif isinstance(obj, type):
+            # Non-protocol class (IntEnum, OffsetResetStrategy, etc.)
+            # Reproduce class with its attributes so the stub doesn't shadow the real definition
+            from enum import IntEnum as _IntEnum
+            bases = [b.__name__ for b in obj.__bases__ if b is not object]
+            base_str = '(%s)' % ', '.join(bases) if bases else ''
+            lines.append('class %s%s:' % (name, base_str))
+            members = {k: v for k, v in obj.__dict__.items()
+                       if not k.startswith('_') and not callable(v)}
+            if members:
+                for k, v in members.items():
+                    # IntEnum members have the enum class as type; use int instead
+                    if issubclass(obj, _IntEnum):
+                        lines.append('    %s: int' % k)
+                    else:
+                        lines.append('    %s: %s' % (k, type(v).__name__))
+            else:
+                lines.append('    ...')
+            lines.append('')
+        else:
+            # Constants
+            type_name = type(obj).__name__
+            lines.append('%s: %s' % (name, type_name))
+            lines.append('')
+
+    return '\n'.join(lines)
+
+
+def generate_all(dry_run=False, check=False):
+    """Generate all stub files. Returns True if all stubs are up to date."""
+    modules = discover_modules()
+    all_up_to_date = True
+
+    for mod_file, (mod, exports) in sorted(modules.items()):
+        content = generate_module(mod, exports)
+        pyi_path = mod_file.replace('.py', '.pyi')
+
+        if dry_run:
+            print('# %s' % pyi_path)
+            print(content)
+            print()
+        elif check:
+            if os.path.exists(pyi_path):
+                with open(pyi_path, 'r') as f:
+                    existing = f.read()
+                if existing != content:
+                    print('OUTDATED: %s' % pyi_path)
+                    all_up_to_date = False
+            else:
+                print('MISSING: %s' % pyi_path)
+                all_up_to_date = False
+        else:
+            with open(pyi_path, 'w') as f:
+                f.write(content)
+            print('Generated: %s' % pyi_path)
+
+    return all_up_to_date
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(description='Generate .pyi stubs for protocol classes')
+    parser.add_argument('--check', action='store_true',
+                        help='Check if stubs are up to date (exit 1 if not)')
+    parser.add_argument('--dry-run', action='store_true',
+                        help='Print to stdout instead of writing files')
+    args = parser.parse_args()
+
+    if args.check:
+        up_to_date = generate_all(check=True)
+        sys.exit(0 if up_to_date else 1)
+    else:
+        generate_all(dry_run=args.dry_run)

kafka/protocol/metadata/__init__.py

@@ -0,0 +1,10 @@
+from .api_versions import *
+from .api_versions import __all__ as api_versions_all
+
+from .find_coordinator import *
+from .find_coordinator import __all__ as find_coordinator_all
+
+from .metadata import *
+from .metadata import __all__ as metadata_all
+
+__all__ = api_versions_all + find_coordinator_all + metadata_all

kafka/protocol/metadata/api_versions.py

@@ -0,0 +1,41 @@
+import io
+
+from ..api_message import ApiMessage
+from ..schemas.fields.codecs import Int16, Int32
+
+
+class ApiVersionsRequest(ApiMessage): pass
+class ApiVersionsResponse(ApiMessage):
+    # ApiVersionsResponse header never uses flexible formats, even if body does
+    @classmethod
+    def parse_header(cls, data, version=None):
+        return cls.header_class.decode(data, flexible=False) # pylint: disable=E1101
+
+    def encode_header(self, flexible=False):
+        return super().encode_header(flexible=False)
+
+    # ApiVersionsResponse body always decodes as version 0 when error is present
+    @classmethod
+    def decode(cls, data, version=None, header=False, framed=False):
+        if not hasattr(data, 'tell'):
+            data = io.BytesIO(data)
+        if framed:
+            size = Int32.decode(data)
+        if header:
+            hdr = cls.parse_header(data)
+        else:
+            hdr = None
+        curr = data.tell()
+        err = Int16.decode(data)
+        data.seek(curr)
+        if err != 0:
+            version = 0
+        ret = super().decode(data, version=version, header=False, framed=False)
+        if hdr is not None:
+            ret._header = hdr
+        return ret
+
+
+__all__ = [
+    'ApiVersionsRequest', 'ApiVersionsResponse',
+]

kafka/protocol/metadata/api_versions.pyi

@@ -0,0 +1,128 @@
+# Generated by generate_stubs.py (Python 3.14)
+import uuid
+from typing import Any, Self
+
+from kafka.protocol.api_message import ApiMessage
+from kafka.protocol.data_container import DataContainer
+
+__all__ = ['ApiVersionsRequest', 'ApiVersionsResponse']
+
+class ApiVersionsRequest(ApiMessage):
+    client_software_name: str
+    client_software_version: str
+    def __init__(
+        self,
+        *args: Any,
+        client_software_name: str = ...,
+        client_software_version: str = ...,
+        version: int | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+    @property
+    def version(self) -> int | None: ...
+    def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+    name: str
+    type: str
+    API_KEY: int
+    API_VERSION: int
+    valid_versions: tuple[int, int]
+    min_version: int
+    max_version: int
+    @property
+    def header(self) -> Any: ...
+    @classmethod
+    def is_request(cls) -> bool: ...
+    def expect_response(self) -> bool: ...
+    def with_header(self, correlation_id: int = 0, client_id: str = "kafka-python") -> None: ...
+
+class ApiVersionsResponse(ApiMessage):
+    class ApiVersion(DataContainer):
+        api_key: int
+        min_version: int
+        max_version: int
+        def __init__(
+            self,
+            *args: Any,
+            api_key: int = ...,
+            min_version: int = ...,
+            max_version: int = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    class SupportedFeatureKey(DataContainer):
+        name: str
+        min_version: int
+        max_version: int
+        def __init__(
+            self,
+            *args: Any,
+            name: str = ...,
+            min_version: int = ...,
+            max_version: int = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    class FinalizedFeatureKey(DataContainer):
+        name: str
+        max_version_level: int
+        min_version_level: int
+        def __init__(
+            self,
+            *args: Any,
+            name: str = ...,
+            max_version_level: int = ...,
+            min_version_level: int = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    error_code: int
+    api_keys: list[ApiVersion]
+    throttle_time_ms: int
+    supported_features: list[SupportedFeatureKey]
+    finalized_features_epoch: int
+    finalized_features: list[FinalizedFeatureKey]
+    zk_migration_ready: bool
+    def __init__(
+        self,
+        *args: Any,
+        error_code: int = ...,
+        api_keys: list[ApiVersion] = ...,
+        throttle_time_ms: int = ...,
+        supported_features: list[SupportedFeatureKey] = ...,
+        finalized_features_epoch: int = ...,
+        finalized_features: list[FinalizedFeatureKey] = ...,
+        zk_migration_ready: bool = ...,
+        version: int | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+    @property
+    def version(self) -> int | None: ...
+    def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+    name: str
+    type: str
+    API_KEY: int
+    API_VERSION: int
+    valid_versions: tuple[int, int]
+    min_version: int
+    max_version: int
+    @property
+    def header(self) -> Any: ...
+    @classmethod
+    def is_request(cls) -> bool: ...
+    def expect_response(self) -> bool: ...
+    def with_header(self, correlation_id: int = 0, client_id: str = "kafka-python") -> None: ...
+    @classmethod
+    def parse_header(cls, data: Any, version: Any = ...) -> Any: ...
+    def encode_header(self, flexible: Any = ...) -> Any: ...

kafka/protocol/metadata/find_coordinator.py

@@ -0,0 +1,19 @@
+from enum import IntEnum
+
+from ..api_message import ApiMessage
+from kafka.util import EnumHelper
+
+
+class CoordinatorType(EnumHelper, IntEnum):
+    GROUP = 0
+    TRANSACTION = 1
+    SHARE = 2
+
+
+class FindCoordinatorRequest(ApiMessage): pass
+class FindCoordinatorResponse(ApiMessage): pass
+
+
+__all__ = [
+    'CoordinatorType', 'FindCoordinatorRequest', 'FindCoordinatorResponse',
+]