Add BitField type to encode/decode at protocol layer

Author: dpkp

kafka/admin/acl_resource.py

@@ -8,7 +8,6 @@
     from kafka.vendor.enum34 import IntEnum
 
 from kafka.errors import IllegalArgumentError
-from kafka.util import from_32_bit_field
 
 
 class ResourceType(IntEnum):
@@ -251,5 +250,5 @@ def validate(self):
             )
 
 
-def valid_acl_operations(int_val):
-     return set([ACLOperation(v) for v in from_32_bit_field(int_val) if v not in (0, 1, 2)])
+def valid_acl_operations(int_vals):
+     return set([ACLOperation(v) for v in int_vals if v not in (0, 1, 2)])

kafka/protocol/admin.py

@@ -8,7 +8,7 @@
     from kafka.vendor.enum34 import IntEnum
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Boolean, Bytes, Int8, Int16, Int32, Int64, Schema, String, Float64, CompactString, CompactArray, TaggedFields
+from kafka.protocol.types import Array, Boolean, Bytes, Int8, Int16, Int32, Int64, Schema, String, Float64, CompactString, CompactArray, TaggedFields, BitField
 
 
 class CreateTopicsResponse_v0(Response):
@@ -338,7 +338,7 @@ class DescribeGroupsResponse_v3(Response):
                 ('client_host', String('utf-8')),
                 ('member_metadata', Bytes),
                 ('member_assignment', Bytes))),
-            ('authorized_operations', Int32)))
+            ('authorized_operations', BitField)))
     )
 
 

kafka/protocol/metadata.py

@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Boolean, Int16, Int32, Schema, String
+from kafka.protocol.types import Array, Boolean, Int16, Int32, Schema, String, BitField
 
 
 class MetadataResponse_v0(Response):
@@ -189,8 +189,8 @@ class MetadataResponse_v8(Response):
                 ('replicas', Array(Int32)),
                 ('isr', Array(Int32)),
                 ('offline_replicas', Array(Int32)))),
-            ('authorized_operations', Int32))),
-        ('authorized_operations', Int32)
+            ('authorized_operations', BitField))),
+        ('authorized_operations', BitField)
     )
 
 

kafka/protocol/types.py

@@ -363,3 +363,34 @@ def decode(self, data):
             return None
         return [self.array_of.decode(data) for _ in range(length)]
 
+
+class BitField(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        return cls.from_32_bit_field(Int32.decode(data))
+
+    @classmethod
+    def encode(cls, vals):
+        # to_32_bit_field returns unsigned val, so we need to
+        # encode >I to avoid crash if/when byte 31 is set
+        # (note that decode as signed still works fine)
+        return struct.Struct('>I').pack(cls.to_32_bit_field(vals))
+
+    @classmethod
+    def to_32_bit_field(cls, vals):
+        value = 0
+        for b in vals:
+            assert 0 <= b < 32
+            value |= 1 << b
+        return value
+
+    @classmethod
+    def from_32_bit_field(cls, value):
+        result = set()
+        count = 0
+        while value != 0:
+            if (value & 1) != 0:
+                result.add(count)
+            count += 1
+            value = (value & 0xFFFFFFFF) >> 1
+        return result

kafka/util.py

@@ -139,20 +139,3 @@ def wrapper(self, *args, **kwargs):
     functools.update_wrapper(wrapper, func)
     return wrapper
 
-
-def to_32_bit_field(vals):
-    value = 0
-    for b in vals:
-        assert 0 <= b < 32
-        value |= 1 << b
-    return value
-
-def from_32_bit_field(value):
-    result = set()
-    count = 0
-    while value != 0:
-        if (value & 1) != 0:
-            result.add(count)
-        count += 1
-        value = (value & 0xFFFFFFFF) >> 1
-    return result

test/test_protocol.py

@@ -2,12 +2,14 @@
 import io
 import struct
 
+import pytest
+
 from kafka.protocol.api import RequestHeader
 from kafka.protocol.fetch import FetchRequest, FetchResponse
 from kafka.protocol.find_coordinator import FindCoordinatorRequest
 from kafka.protocol.message import Message, MessageSet, PartialMessage
 from kafka.protocol.metadata import MetadataRequest
-from kafka.protocol.types import Int16, Int32, Int64, String, UnsignedVarInt32, CompactString, CompactArray, CompactBytes
+from kafka.protocol.types import Int16, Int32, Int64, String, UnsignedVarInt32, CompactString, CompactArray, CompactBytes, BitField
 
 
 def test_create_message():
@@ -332,3 +334,11 @@ def test_compact_data_structs():
     assert CompactBytes.decode(io.BytesIO(b'\x01')) == b''
     enc = CompactBytes.encode(b'foo')
     assert CompactBytes.decode(io.BytesIO(enc)) == b'foo'
+
+
+@pytest.mark.parametrize(('test_set',), [
+    (set([0, 1, 5, 10, 31]),),
+    (set(range(32)),),
+])
+def test_bit_field(test_set):
+    assert BitField.decode(io.BytesIO(BitField.encode(test_set))) == test_set

test/test_util.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from kafka.util import ensure_valid_topic_name, from_32_bit_field, to_32_bit_field
+from kafka.util import ensure_valid_topic_name
 
 @pytest.mark.parametrize(('topic_name', 'expectation'), [
     (0, pytest.raises(TypeError)),
@@ -23,10 +23,3 @@ def test_topic_name_validation(topic_name, expectation):
     with expectation:
         ensure_valid_topic_name(topic_name)
 
-
-@pytest.mark.parametrize(('test_set',), [
-    (set([0, 1, 5, 10, 31]),),
-    (set(range(32)),),
-])
-def test_32_bit_field(test_set):
-    assert from_32_bit_field(to_32_bit_field(test_set)) == test_set