Separate ApiStructArray

Author: dpkp

kafka/protocol/new/__init__.py

@@ -1,7 +1,8 @@
 # These need to get registered before ApiMessage metaclass starts building classes
 from .field_basic import FieldBasicType
-from .api_array import ApiArray
 from .api_struct import ApiStruct
+from .api_array import ApiArray
+from .api_struct_array import ApiStructArray
 
 from .api_message import ApiMessage
 from .api_header import ApiHeader, RequestHeader, ResponseHeader

kafka/protocol/new/api_array.py

@@ -1,4 +1,5 @@
 from .field import Field
+from .field_basic import FieldBasicType
 from ..types import (
     Array, CompactArray,
     UnsignedVarInt32, Int32,
@@ -7,44 +8,37 @@
 
 class ApiArray(Field):
     @classmethod
-    def parse_json(cls, json):
+    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 FieldBasicType.parse_json(inner_json)
+
+    @classmethod
+    def parse_array_type(cls, json):
         if json['type'].startswith('[]'):
-            inner_type_str = json['type'][2:]
-            if inner_type_str.startswith('[]'): # this would be strange...
-                return None
-            inner_json = {**json, 'type': inner_type_str}
-            inner_type = super().parse_json(inner_json)
-            if inner_type is not None:
-                return cls(json, array_of=inner_type)
+            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 # Field (ApiStruct or FieldBasicType)
+        self.array_of = array_of # FieldBasicType
 
     def is_array(self):
         return True
 
-    def is_struct_array(self):
-        return self.array_of.is_struct()
-
-    @property
-    def fields(self):
-        if self.is_struct_array():
-            return self.array_of.fields
-
-    def has_data_class(self):
-        return self.is_struct_array()
-
-    @property
-    def data_class(self):
-        if self.has_data_class():
-            return self.array_of.data_class
-        else:
-            raise ValueError('Non-struct field does not have a data_class!')
-
-    def __call__(self, *args, **kw):
-        return self.data_class(*args, **kw) # pylint: disable=E1102
-
     def to_schema(self, version, compact=False, tagged=False):
         if not self.for_version_q(version) or self.tagged_field_q(version):
             return None

kafka/protocol/new/api_struct_array.py

@@ -0,0 +1,47 @@
+from .api_array import ApiArray
+from .api_struct import ApiStruct
+
+
+class ApiStructArray(ApiArray):
+    @classmethod
+    def parse_inner_type(cls, json):
+        # ApiStructArray requires non-empty fields
+        if 'fields' not in json:
+            return
+        assert len(json['fields']) > 0, 'Unexpected empty fields in json'
+        type_str = cls.parse_array_type(json)
+        if type_str is None:
+            return
+        inner_json = {**json, 'type': type_str}
+        return ApiStruct.parse_json(inner_json)
+
+    @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 StructArray!'
+        super().__init__(json, array_of=array_of)
+        # map_key will be (idx, field) of the mapKey field if found
+        self.map_key = next(lambda i, field: field._json.get('mapKey'), enumerate(self._fields), None)
+
+    def is_struct_array(self):
+        return True
+
+    @property
+    def fields(self):
+        return self.array_of.fields
+
+    def has_data_class(self):
+        return True
+
+    @property
+    def data_class(self):
+        return self.array_of.data_class
+
+    def __call__(self, *args, **kw):
+        return self.data_class(*args, **kw) # pylint: disable=E1102