diff --git a/kafka/protocol/new/api_message.py b/kafka/protocol/new/api_message.py index 97af532d1..6e4a689d2 100644 --- a/kafka/protocol/new/api_message.py +++ b/kafka/protocol/new/api_message.py @@ -86,15 +86,8 @@ def __init_subclass__(cls, **kw): def __init__(self, *args, **kwargs): self._header = None - self._version = None - if 'version' in kwargs: - self.API_VERSION = kwargs['version'] - if len(args) > 0: - untagged_fields = self._struct.untagged_fields(self.API_VERSION) - if len(args) != len(untagged_fields): - raise RuntimeError('Unable to init ApiMessage via positional args: unexpected len') - kwargs.update({field.name: args[i] for i, field in enumerate(untagged_fields)}) - args = () + if 'version' not in kwargs: + kwargs['version'] = self._class_version # pylint: disable=E1101 super().__init__(*args, **kwargs) @classproperty diff --git a/kafka/protocol/new/data_container.py b/kafka/protocol/new/data_container.py index f2cba89b1..6df3ce8dc 100644 --- a/kafka/protocol/new/data_container.py +++ b/kafka/protocol/new/data_container.py @@ -22,9 +22,19 @@ def __init_subclass__(cls, **kwargs): field.set_data_class(type(field.type_str, (DataContainer,), {'_struct': field})) setattr(cls, field.type_str, field.data_class) - def __init__(self, version=None, **field_vals): + def __init__(self, *args, version=None, **field_vals): assert self._struct is not None self._version = version + # Support positional arg init for convenience + if len(args) > 0: + if self._version is not None: + field_args = [field for field in self._struct._fields if field.for_version_q(self._version)] + else: + field_args = self._struct._fields + if len(args) > len(field_args): + raise RuntimeError('Unable to init DataContainer with %d positional args: unexpected %d' % (len(args), len(field_args))) + field_vals.update({field_args[i].name: arg for i, arg in enumerate(args)}) + args = () self.tags = None self.unknown_tags = None for field in self._struct._fields: