kafka-python 3.0.0__py3-none-any.whl

kafka/protocol/sasl.pyi

@@ -0,0 +1,126 @@
+# 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.api_data import ApiData
+
+__all__ = ['SaslHandshakeRequest', 'SaslHandshakeResponse', 'SaslAuthenticateRequest', 'SaslAuthenticateResponse', 'SaslBytesRequest', 'SaslBytesResponse']
+
+class SaslHandshakeRequest(ApiMessage):
+    mechanism: str
+    def __init__(
+        self,
+        *args: Any,
+        mechanism: 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 SaslHandshakeResponse(ApiMessage):
+    error_code: int
+    mechanisms: list[str]
+    def __init__(
+        self,
+        *args: Any,
+        error_code: int = ...,
+        mechanisms: list[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 SaslAuthenticateRequest(ApiMessage):
+    auth_bytes: bytes | ApiData
+    def __init__(
+        self,
+        *args: Any,
+        auth_bytes: bytes | ApiData = ...,
+        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 SaslAuthenticateResponse(ApiMessage):
+    error_code: int
+    error_message: str | None
+    auth_bytes: bytes | ApiData
+    session_lifetime_ms: int
+    def __init__(
+        self,
+        *args: Any,
+        error_code: int = ...,
+        error_message: str | None = ...,
+        auth_bytes: bytes | ApiData = ...,
+        session_lifetime_ms: 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: ...
+    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 SaslBytesRequest:
+    API_VERSION: int
+
+class SaslBytesResponse:
+    ...

kafka/protocol/schemas/__init__.py

@@ -0,0 +1,7 @@
+from .fields import *
+from .fields import __all__ as fields_all
+
+from .load_json import load_json
+
+
+__all__ = fields_all + ['load_json']

kafka/protocol/schemas/fields/__init__.py

@@ -0,0 +1,7 @@
+from .base import BaseField
+from .simple import SimpleField
+from .struct import StructField
+from .array import ArrayField
+from .struct_array import StructArrayField
+
+__all__ = ["BaseField", "SimpleField", "StructField", "ArrayField", "StructArrayField"]

kafka/protocol/schemas/fields/array.py

@@ -0,0 +1,127 @@
+from .base import BaseField
+from .simple import SimpleField
+from .codecs import (
+    UnsignedVarInt32, Int32,
+)
+
+
+class ArrayField(BaseField):
+    @classmethod
+    def parse_inner_type(cls, json):
+        if 'fields' in json:
+            return
+        type_str = cls.parse_array_type(json)
+        if type_str is not None:
+            inner_json = {**json, 'type': type_str}
+            return SimpleField.parse_json(inner_json)
+
+    @classmethod
+    def parse_array_type(cls, json):
+        if json['type'].startswith('[]'):
+            type_str = json['type'][2:]
+            assert not type_str.startswith('[]'), 'Unexpected double-array type: %s' % json['type']
+            return type_str
+
+    @classmethod
+    def parse_json(cls, json):
+        inner_type = cls.parse_inner_type(json)
+        if inner_type is not None:
+            return cls(json, array_of=inner_type)
+
+    def __init__(self, json, array_of=None):
+        if array_of is None:
+            array_of = self.parse_inner_type(json)
+            assert array_of is not None, 'json does not contain a (simple) Array!'
+        super().__init__(json)
+        self.array_of = array_of # SimpleField
+
+    def is_array(self):
+        return True
+
+    def _calculate_default(self, default):
+        if default == 'null':
+            return None
+        elif not default:
+            return []
+        else:
+            raise ValueError('Invalid default for field %s. The only valid default is empty or null.' % self._name)
+
+    def encode(self, items, version=None, compact=False, tagged=False):
+        if compact:
+            size = UnsignedVarInt32.encode(len(items) + 1 if items is not None else 0)
+        else:
+            size = Int32.encode(len(items) if items is not None else -1)
+        if items is None:
+            return size
+        fields = [self.array_of.encode(item, version=version, compact=compact, tagged=tagged)
+                  for item in items]
+        return b''.join([size] + fields)
+
+    def encode_into(self, items, out, version=None, compact=False, tagged=False):
+        if compact:
+            UnsignedVarInt32.encode_into(out, len(items) + 1 if items is not None else 0)
+        else:
+            Int32.encode_into(out, len(items) if items is not None else -1)
+        if items is None:
+            return
+        encode_into = self.array_of.encode_into
+        for item in items:
+            encode_into(item, out, version=version, compact=compact, tagged=tagged)
+
+    def emit_encode_into(self, ctx, val_expr, indent, version=None, compact=False, tagged=False):
+        if compact:
+            an = ctx.next_var('an')
+            ctx.emit(indent, '%s = len(%s) + 1 if %s is not None else 0' % (an, val_expr, val_expr))
+            UnsignedVarInt32.emit_encode_into(ctx, an, indent)
+            ctx.emit(indent, 'if %s is not None:' % val_expr)
+        else:
+            ctx.emit(indent, 'if %s is None:' % val_expr)
+            ctx.emit(indent, "    pack_into('>i', buf, pos, -1)")
+            ctx.emit(indent, '    pos += 4')
+            ctx.emit(indent, 'else:')
+            ctx.emit(indent, "    pack_into('>i', buf, pos, len(%s))" % val_expr)
+            ctx.emit(indent, '    pos += 4')
+        guard = indent + '    '
+        item_var = ctx.next_var('ai')
+        ctx.emit(guard, 'for %s in %s:' % (item_var, val_expr))
+        self.array_of.emit_encode_into(ctx, item_var, guard + '    ',
+                                       version=version, compact=compact, tagged=tagged)
+
+    def emit_decode_from(self, ctx, var_name, indent, version=None, compact=False, tagged=False):
+        n = ctx.next_var('n')
+        if compact:
+            UnsignedVarInt32.emit_decode_from(ctx, n, indent)
+            ctx.emit(indent, '%s -= 1' % n)
+        else:
+            ctx.emit(indent, '%s = unpack_from(">i", data, pos)[0]' % n)
+            ctx.emit(indent, 'pos += 4')
+        ctx.emit(indent, 'if %s == -1:' % n)
+        ctx.emit(indent, '    %s = None' % var_name)
+        ctx.emit(indent, 'else:')
+        inner_indent = indent + '    '
+        ctx.emit(inner_indent, '%s = []' % var_name)
+        idx = ctx.next_var('idx')
+        item = ctx.next_var('item')
+        ctx.emit(inner_indent, 'for %s in range(%s):' % (idx, n))
+        self.array_of.emit_decode_from(ctx, item, inner_indent + '    ',
+                                        version=version, compact=compact, tagged=tagged)
+        ctx.emit(inner_indent, '    %s.append(%s)' % (var_name, item))
+
+    def decode(self, data, version=None, compact=False, tagged=False):
+        if compact:
+            size = UnsignedVarInt32.decode(data)
+            size -= 1
+        else:
+            size = Int32.decode(data)
+        if size == -1:
+            return None
+        return [self.array_of.decode(data, version=version, compact=compact, tagged=tagged)
+                for _ in range(size)]
+
+    def to_json(self, val):
+        if val is None:
+            return None
+        return [self.array_of.to_json(i) for i in val]
+
+    def __repr__(self):
+        return 'ArrayField(%s)' % self._json

kafka/protocol/schemas/fields/base.py

@@ -0,0 +1,156 @@
+import re
+
+
+class BaseField:
+    FIELD_TYPES = []
+
+    def __init_subclass__(cls, **kw):
+        cls.FIELD_TYPES.append(cls)
+
+    @classmethod
+    def parse_json_fields(cls, json):
+        return tuple(map(BaseField.parse_json, json.get('fields', []))) or None # Note: DFS Field construction
+
+    @classmethod
+    def parse_json(cls, json):
+        if 'type' not in json:
+            raise ValueError('No type found in json')
+        type_str = json['type']
+        for field_type in cls.FIELD_TYPES:
+            maybe_type = field_type.parse_json(json)
+            if maybe_type is not None:
+                return maybe_type
+        else:
+            raise ValueError('Unable to parse field type: %s' % type_str)
+
+    @classmethod
+    def underscore_name(cls, name):
+        return re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()
+
+    @classmethod
+    def parse_versions(cls, versions):
+        if versions is None or versions.strip() == '':
+            return None
+        elif versions.strip() == 'none':
+            return (-1, -1)
+        elif versions[-1] == '+':
+            return (int(versions[:-1]), 32767)
+        elif versions.find('-') == -1:
+            return (int(versions), int(versions))
+        else:
+            return tuple(map(int, versions.split('-')))
+
+    def __init__(self, json):
+        self._json = json
+        self._name = json['name']
+        self.name = self.underscore_name(self._name)
+        # versions tells when to include the field in encode / decode
+        # 'validVersions' is included for top-level request/response structs
+        # otherwise this should always be 'versions'
+        for versions_key in ('versions', 'validVersions'):
+            if versions_key in json:
+                self._versions = self.parse_versions(json[versions_key])
+                if self._versions is not None:
+                    break
+        else:
+            raise ValueError('Unable to find versions in json!')
+        # If the field is tagged, it should have 'tag' and 'taggedVersions'
+        self._tag = json.get('tag')
+        self._tagged_versions = self.parse_versions(json.get('taggedVersions'))
+        # flexibleVersions is used to override 'compact' string/bytes encoding for a specific field
+        # currently only used in RequestHeader client_id
+        self._flexible_versions = self.parse_versions(json.get('flexibleVersions'))
+        self._nullable_versions = self.parse_versions(json.get('nullableVersions'))
+        self._type_str = json['type']
+        if 'fields' in json:
+            self._fields = self.parse_json_fields(json)
+        else:
+            self._fields = None
+        self._ignorable = json.get('ignorable')
+        self._entity_type = json.get('entityType')
+        self._about = json.get('about', '')
+        self._zero_copy = json.get('zeroCopy') # ?
+
+    @property
+    def type_str(self):
+        return self._type_str
+
+    @property
+    def tag(self):
+        return self._tag
+
+    # Override in subclasses
+    def is_array(self):
+        return False
+
+    def is_struct(self):
+        return False
+
+    def is_struct_array(self):
+        return False
+
+    def is_batchable(self):
+        return False
+
+    def has_data_class(self):
+        return False
+
+    def __call__(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @property
+    def fields(self):
+        return self._fields
+
+    @property
+    def min_version(self):
+        return self._versions[0]
+
+    @property
+    def max_version(self):
+        return self._versions[1]
+
+    def for_version_q(self, version):
+        return self._versions[0] <= version <= self._versions[1]
+
+    def nullable_for_version_q(self, version):
+        if self._nullable_versions is None:
+            return False
+        return self._nullable_versions[0] <= version <= self._nullable_versions[1]
+
+    def tagged_field_q(self, version):
+        if self._tag is None or self._tagged_versions is None:
+            return False
+        elif not self._tagged_versions[0] <= version <= self._tagged_versions[1]:
+            return False
+        else:
+            return True
+
+    def _calculate_default(self, default):
+        return default
+
+    @property
+    def default(self):
+        try:
+            return self._default # pylint: disable=E1101
+        except AttributeError:
+            self._default = self._calculate_default(self._json.get('default', ''))
+            return self._default
+
+    def encode(self, value, version=None, compact=False, tagged=False):
+        raise NotImplementedError
+
+    def decode(self, data, version=None, compact=False, tagged=False):
+        raise NotImplementedError
+
+    def __repr__(self):
+        return 'BaseField(%s)' % self._json
+
+    def __eq__(self, other):
+        if self.__class__ != other.__class__:
+            return False
+        if self.name != other.name:
+            return False
+        if self.tag != other.tag:
+            return False
+        return True

kafka/protocol/schemas/fields/codecs/__init__.py

@@ -0,0 +1,12 @@
+from .encode_buffer import EncodeBuffer, EncodeBufferPool
+from .tagged_fields import TaggedFields
+from .types import (
+    BitField, Boolean, UUID, Bytes, String,
+    Int8, Int16, Int32, Int64, UnsignedInt16, UnsignedVarInt32, Float64,
+)
+
+__all__ = [
+    'BitField', 'Boolean', 'UUID', 'Bytes', 'String',
+    'Int8', 'Int16', 'Int32', 'Int64', 'UnsignedInt16', 'UnsignedVarInt32', 'Float64',
+    'TaggedFields', 'EncodeBuffer',
+]

kafka/protocol/schemas/fields/codecs/encode_buffer.py

@@ -0,0 +1,82 @@
+import threading
+
+
+class EncodeBuffer:
+    """Growable buffer for encode_into operations."""
+    __slots__ = ('buf', 'pos')
+
+    def __init__(self, size=65536):
+        self.buf = bytearray(size)
+        self.pos = 0
+
+    def reset(self):
+        """Reset position to 0 for reuse. The buffer retains its current size."""
+        self.pos = 0
+
+    def ensure(self, needed):
+        if self.pos + needed > len(self.buf):
+            new_size = max(len(self.buf) * 2, self.pos + needed)
+            new_buf = bytearray(new_size)
+            new_buf[:self.pos] = self.buf[:self.pos]
+            self.buf = new_buf
+
+    def result(self):
+        # Return a bytearray slice (one copy) rather than bytes(self.buf[:pos])
+        # (two copies - the slice creates a bytearray, then bytes() copies
+        # again). Downstream consumers (protocol codec slice assignment,
+        # socket.send) accept bytearray transparently.
+        return self.buf[:self.pos]
+
+
+class EncodeBufferPool:
+    """Thread-local pool of reusable EncodeBuffer objects.
+
+    Each thread gets its own buffer that grows to match the largest message
+    encoded on that thread and stays that size - avoiding repeated allocation
+    of large bytearrays.
+
+    Usage:
+        with EncodeBufferPool.acquire() as out:
+            fast_encode(item, out)
+            return out.result()
+    """
+    _local = threading.local()
+
+    @classmethod
+    def acquire(cls):
+        """Return a context manager that provides a reset EncodeBuffer."""
+        return _PooledBuffer(cls)
+
+    @classmethod
+    def _get(cls):
+        buf = getattr(cls._local, 'buf', None)
+        if buf is None:
+            buf = EncodeBuffer()
+            cls._local.buf = buf
+        buf.reset()
+        return buf
+
+    @classmethod
+    def _release(cls, buf):
+        # Nothing to do - the buffer stays on the thread-local.
+        # Future: could cap max size to prevent memory leaks from
+        # one-off large messages.
+        pass
+
+
+class _PooledBuffer:
+    """Context manager for EncodeBufferPool.acquire()."""
+    __slots__ = ('_pool', '_buf')
+
+    def __init__(self, pool):
+        self._pool = pool
+        self._buf = None
+
+    def __enter__(self):
+        self._buf = self._pool._get()
+        return self._buf
+
+    def __exit__(self, *exc):
+        self._pool._release(self._buf)
+        self._buf = None
+        return False

kafka/protocol/schemas/fields/codecs/tagged_fields.py

@@ -0,0 +1,109 @@
+from .types import UnsignedVarInt32
+
+
+def _uvarint_size(value):
+    """Return the number of bytes needed to encode value as an unsigned varint."""
+    if value < 0x80:
+        return 1
+    elif value < 0x4000:
+        return 2
+    elif value < 0x200000:
+        return 3
+    elif value < 0x10000000:
+        return 4
+    else:
+        return 5
+
+
+class TaggedFields:
+    def __init__(self, fields):
+        self._fields = list(fields) # listify to clean filter() inputs
+        self._tags = {field.tag: field for field in self._fields}
+        self._names = {field.name: field for field in self._fields}
+
+    def encode(self, item, version=None):
+        if isinstance(item, dict):
+            tags = [(self._names[name].tag, val)
+                    for name, val in item.items()
+                    if name in self._names
+                    and val != self._names[name].default]
+        else:
+            tags = [(self._names[name].tag, getattr(item, name))
+                    for name in self._names
+                    if hasattr(item, name)
+                    and getattr(item, name) != self._names[name].default]
+        ret = [UnsignedVarInt32.encode(len(tags))]
+        for tag, val in tags:
+            ret.append(UnsignedVarInt32.encode(tag))
+            # struct tags have an empty nested tagged fields (tagged=None)
+            encoded_val = self._tags[tag].encode(val, version=version, compact=True, tagged=None)
+            ret.append(UnsignedVarInt32.encode(len(encoded_val)))
+            ret.append(encoded_val)
+        return b''.join(ret)
+
+    def encode_into(self, item, out, version=None):
+        if isinstance(item, dict):
+            tags = [(self._names[name].tag, name)
+                    for name, val in item.items()
+                    if name in self._names
+                    and val != self._names[name].default]
+        else:
+            tags = [(self._names[name].tag, name)
+                    for name in self._names
+                    if hasattr(item, name)
+                    and getattr(item, name) != self._names[name].default]
+        UnsignedVarInt32.encode_into(out, len(tags))
+        for tag, name in tags:
+            field = self._tags[tag]
+            val = item.get(name) if isinstance(item, dict) else getattr(item, name)
+            UnsignedVarInt32.encode_into(out, tag)
+            # Encode value into buffer after reserving space for size prefix.
+            # Strategy: assume size fits in 1 byte (< 128), encode value,
+            # then fix up if the size varint is larger.
+            size_pos = out.pos
+            out.pos += 1  # reserve 1 byte for size
+            val_start = out.pos
+            # struct tags have an empty nested tagged fields (tagged=None)
+            field.encode_into(val, out, version=version, compact=True, tagged=None)
+            val_size = out.pos - val_start
+            actual_size_len = _uvarint_size(val_size)
+            if actual_size_len == 1:
+                # Common case: size fits in 1 byte, just write it
+                out.buf[size_pos] = val_size
+            else:
+                # Rare case: size needs more bytes. Shift the value data forward.
+                extra = actual_size_len - 1
+                out.ensure(extra)
+                # Move value bytes forward to make room
+                out.buf[val_start + extra:out.pos + extra] = out.buf[val_start:out.pos]
+                out.pos += extra
+                # Write the multi-byte size at size_pos
+                saved_pos = out.pos
+                out.pos = size_pos
+                UnsignedVarInt32.encode_into(out, val_size)
+                out.pos = saved_pos
+
+    def decode(self, data, version=None):
+        num_fields = UnsignedVarInt32.decode(data)
+        ret = {}
+        for i in range(num_fields):
+            tag = UnsignedVarInt32.decode(data)
+            size = UnsignedVarInt32.decode(data)
+            if tag in self._tags:
+                field = self._tags[tag]
+                # struct tags have an empty nested tagged fields (tagged=None)
+                ret[field.name] = field.decode(data, version=version, compact=True, tagged=None)
+            else:
+                ret['_%d' % tag] = data.read(size)
+        return ret
+
+    @classmethod
+    def decode_empty(cls, data):
+        assert UnsignedVarInt32.decode(data) == 0
+
+    @classmethod
+    def encode_empty(cls):
+        return UnsignedVarInt32.encode(0)
+
+    def __repr__(self):
+        return 'TaggedFields(%s)' % list(self._names.keys())