protobuf API¶

def encode_varint(value: int) -> bytes:
    """Encode an unsigned integer as a varint.

    Args:
        value: Non-negative integer to encode.

    Returns:
        Varint-encoded bytes.
    """
    if value < 0:
        # Proto3 treats negative int32/int64 as 10-byte two's complement
        value = value & 0xFFFFFFFFFFFFFFFF
    # Fast-paths for common small values (tags, lengths, small ints)
    if value < 0x80:
        return bytes((value,))
    if value < 0x4000:
        return bytes(((value & 0x7F) | 0x80, value >> 7))
    buf = bytearray()
    while value > 0x7F:
        buf.append((value & 0x7F) | 0x80)
        value >>= 7
    buf.append(value & 0x7F)
    return bytes(buf)

def decode_varint(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int]:
    """Decode a varint from *data* starting at *pos*.

    Args:
        data: Buffer to read from.
        pos: Start offset.

    Returns:
        Tuple of (decoded value, new position after varint).

    Raises:
        ValueError: If the varint is malformed or exceeds 10 bytes.
    """
    result = 0
    shift = 0
    while True:
        if pos >= len(data):
            raise ValueError("Unexpected end of data while reading varint")
        byte = data[pos]
        result |= (byte & 0x7F) << shift
        pos += 1
        if not (byte & 0x80):
            break
        shift += 7
        if shift >= 70:
            raise ValueError("Varint too long (> 10 bytes)")
    return result, pos

def zigzag_encode(value: int) -> int:
    """ZigZag-encode a signed integer.

    Maps signed integers to unsigned: 0→0, -1→1, 1→2, -2→3, …

    Args:
        value: Signed integer.

    Returns:
        ZigZag-encoded unsigned integer.
    """
    return (value << 1) ^ (value >> 63)

def zigzag_decode(value: int) -> int:
    """ZigZag-decode an unsigned integer back to signed.

    Args:
        value: ZigZag-encoded unsigned integer.

    Returns:
        Original signed integer.
    """
    return (value >> 1) ^ -(value & 1)

def make_tag(field_number: int, wire_type: int) -> bytes:
    """Pack a field number and wire type into a tag varint.

    Args:
        field_number: Proto field number (1–536870911).
        wire_type: Wire type (0–5).

    Returns:
        Varint-encoded tag bytes.
    """
    key = (field_number, wire_type)
    cached = _TAG_CACHE.get(key)
    if cached is not None:
        return cached
    tag = encode_varint((field_number << 3) | wire_type)
    _TAG_CACHE[key] = tag
    return tag

def decode_tag(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int, int]:
    """Decode a tag varint into field number and wire type.

    Args:
        data: Buffer.
        pos: Start offset.

    Returns:
        Tuple of (field_number, wire_type, new position).
    """
    tag, pos = decode_varint(data, pos)
    return tag >> 3, tag & 0x07, pos

def encode_fixed32(value: int) -> bytes:
    """Encode a 32-bit value in little-endian."""
    return struct.pack("<I", value & 0xFFFFFFFF)

def decode_fixed32(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int]:
    """Decode a 32-bit little-endian unsigned integer."""
    if pos + 4 > len(data):
        raise ValueError("Unexpected end of data reading fixed32")
    return struct.unpack_from("<I", data, pos)[0], pos + 4

def encode_sfixed32(value: int) -> bytes:
    """Encode a signed 32-bit value in little-endian."""
    return struct.pack("<i", value)

def decode_sfixed32(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int]:
    """Decode a signed 32-bit little-endian integer."""
    if pos + 4 > len(data):
        raise ValueError("Unexpected end of data reading sfixed32")
    return struct.unpack_from("<i", data, pos)[0], pos + 4

def encode_fixed64(value: int) -> bytes:
    """Encode a 64-bit value in little-endian."""
    return struct.pack("<Q", value & 0xFFFFFFFFFFFFFFFF)

def decode_fixed64(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int]:
    """Decode a 64-bit little-endian unsigned integer."""
    if pos + 8 > len(data):
        raise ValueError("Unexpected end of data reading fixed64")
    return struct.unpack_from("<Q", data, pos)[0], pos + 8

def encode_sfixed64(value: int) -> bytes:
    """Encode a signed 64-bit value in little-endian."""
    return struct.pack("<q", value)

def decode_sfixed64(data: bytes | bytearray | memoryview, pos: int) -> tuple[int, int]:
    """Decode a signed 64-bit little-endian integer."""
    if pos + 8 > len(data):
        raise ValueError("Unexpected end of data reading sfixed64")
    return struct.unpack_from("<q", data, pos)[0], pos + 8

def encode_float(value: float) -> bytes:
    """Encode a 32-bit float."""
    return struct.pack("<f", value)

def decode_float(data: bytes | bytearray | memoryview, pos: int) -> tuple[float, int]:
    """Decode a 32-bit float."""
    if pos + 4 > len(data):
        raise ValueError("Unexpected end of data reading float")
    return struct.unpack_from("<f", data, pos)[0], pos + 4

def encode_double(value: float) -> bytes:
    """Encode a 64-bit double."""
    return struct.pack("<d", value)

def decode_double(data: bytes | bytearray | memoryview, pos: int) -> tuple[float, int]:
    """Decode a 64-bit double."""
    if pos + 8 > len(data):
        raise ValueError("Unexpected end of data reading double")
    return struct.unpack_from("<d", data, pos)[0], pos + 8

def oneof(group_name: str) -> OneofGroup:
    """Create a oneof group marker for use in field metadata.

    Usage::

        @message
        class Msg:
            text: str = field(1, oneof="body")
            image: bytes = field(2, oneof="body")

    Args:
        group_name: Name of the oneof group.

    Returns:
        OneofGroup marker.
    """
    return OneofGroup(group_name)

def field(
    number: int,
    *,
    default: Any = dataclasses.MISSING,
    default_factory: Any = dataclasses.MISSING,
    oneof: str | None = None,
) -> Any:
    """Define a proto field with its field number.

    Args:
        number: Proto field number (must be >= 1).
        default: Default value for the field.
        default_factory: Factory for mutable default values.
        oneof: Optional oneof group name.

    Returns:
        A ``dataclasses.Field`` with proto metadata.
    """
    if number < 1:
        raise ValueError(f"Field number must be >= 1, got {number}")
    metadata: dict[str, Any] = {"proto_number": number}
    if oneof is not None:
        metadata["proto_oneof"] = oneof
    kwargs: dict[str, Any] = {"metadata": metadata}
    if default is not dataclasses.MISSING:
        kwargs["default"] = default
    elif default_factory is not dataclasses.MISSING:
        kwargs["default_factory"] = default_factory
    return dataclasses.field(**kwargs)

def message(cls: type) -> type:
    """Decorator that turns a class into a proto3 message.

    Applies ``@dataclass`` (if not already applied) and injects proto3
    ``serialize()``, ``parse()``, ``to_dict()``, ``from_dict()`` methods.

    Usage::

        @message
        class Person:
            name: str = field(1)
            id: int32 = field(2)

    Args:
        cls: The class to decorate.

    Returns:
        The decorated class with proto3 capabilities.
    """
    if not dataclasses.is_dataclass(cls):
        _apply_proto3_defaults(cls)
        cls = dataclasses.dataclass(cls)

    # Build descriptor
    descriptor = _MessageDescriptor(cls)
    descriptor._bind_handlers()
    cls._proto_descriptor = descriptor  # type: ignore[attr-defined]

    # Inject methods
    cls.serialize = _msg_serialize  # type: ignore[attr-defined]
    cls.parse = classmethod(_msg_parse)  # type: ignore[attr-defined]
    cls.to_dict = _msg_to_dict  # type: ignore[attr-defined]
    cls.from_dict = classmethod(_msg_from_dict)  # type: ignore[attr-defined]
    cls.byte_size = _msg_byte_size  # type: ignore[attr-defined]

    # Add _unknown_fields support
    original_init = cls.__init__

    def _new_init(self: Any, *args: Any, **kwargs: Any) -> None:
        original_init(self, *args, **kwargs)
        if not hasattr(self, "_unknown_fields"):
            object.__setattr__(self, "_unknown_fields", [])

    cls.__init__ = _new_init  # type: ignore[attr-defined]

    return cls