summaryrefslogtreecommitdiff
path: root/kafka/record/util.py
blob: 3b712005de748426513e36b306f9ac5b0ca4adb7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import binascii

from kafka.record._crc32c import crc as crc32c_py
try:
    from crc32c import crc32c as crc32c_c
except ImportError:
    crc32c_c = None


def encode_varint(value, write):
    """ Encode an integer to a varint presentation. See
    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
    on how those can be produced.

        Arguments:
            value (int): Value to encode
            write (function): Called per byte that needs to be writen

        Returns:
            int: Number of bytes written
    """
    value = (value << 1) ^ (value >> 63)

    if value <= 0x7f:  # 1 byte
        write(value)
        return 1
    if value <= 0x3fff:  # 2 bytes
        write(0x80 | (value & 0x7f))
        write(value >> 7)
        return 2
    if value <= 0x1fffff:  # 3 bytes
        write(0x80 | (value & 0x7f))
        write(0x80 | ((value >> 7) & 0x7f))
        write(value >> 14)
        return 3
    if value <= 0xfffffff:  # 4 bytes
        write(0x80 | (value & 0x7f))
        write(0x80 | ((value >> 7) & 0x7f))
        write(0x80 | ((value >> 14) & 0x7f))
        write(value >> 21)
        return 4
    if value <= 0x7ffffffff:  # 5 bytes
        write(0x80 | (value & 0x7f))
        write(0x80 | ((value >> 7) & 0x7f))
        write(0x80 | ((value >> 14) & 0x7f))
        write(0x80 | ((value >> 21) & 0x7f))
        write(value >> 28)
        return 5
    else:
        # Return to general algorithm
        bits = value & 0x7f
        value >>= 7
        i = 0
        while value:
            write(0x80 | bits)
            bits = value & 0x7f
            value >>= 7
            i += 1
    write(bits)
    return i


def size_of_varint(value):
    """ Number of bytes needed to encode an integer in variable-length format.
    """
    value = (value << 1) ^ (value >> 63)
    if value <= 0x7f:
        return 1
    if value <= 0x3fff:
        return 2
    if value <= 0x1fffff:
        return 3
    if value <= 0xfffffff:
        return 4
    if value <= 0x7ffffffff:
        return 5
    if value <= 0x3ffffffffff:
        return 6
    if value <= 0x1ffffffffffff:
        return 7
    if value <= 0xffffffffffffff:
        return 8
    if value <= 0x7fffffffffffffff:
        return 9
    return 10


def decode_varint(buffer, pos=0):
    """ Decode an integer from a varint presentation. See
    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
    on how those can be produced.

        Arguments:
            buffer (bytearray): buffer to read from.
            pos (int): optional position to read from

        Returns:
            (int, int): Decoded int value and next read position
    """
    result = buffer[pos]
    if not (result & 0x81):
        return (result >> 1), pos + 1
    if not (result & 0x80):
        return (result >> 1) ^ (~0), pos + 1

    result &= 0x7f
    pos += 1
    shift = 7
    while 1:
        b = buffer[pos]
        result |= ((b & 0x7f) << shift)
        pos += 1
        if not (b & 0x80):
            return ((result >> 1) ^ -(result & 1), pos)
        shift += 7
        if shift >= 64:
            raise ValueError("Out of int64 range")


_crc32c = crc32c_py
if crc32c_c is not None:
    _crc32c = crc32c_c


def calc_crc32c(memview, _crc32c=_crc32c):
    """ Calculate CRC-32C (Castagnoli) checksum over a memoryview of data
    """
    return _crc32c(memview)


def calc_crc32(memview):
    """ Calculate simple CRC-32 checksum over a memoryview of data
    """
    crc = binascii.crc32(memview) & 0xffffffff
    return crc