def encode_varint(value):
"""Encode a long or int into a bytearray."""
output = bytearray()
if value > (2**63) or value < -(2**63):
raise EncoderException("Value %d above maximum varint size" % value)
try:
encoder._EncodeSignedVarint(_gen_append_bytearray(output), value)
except (struct.error, ValueError) as exc:
six.raise_from(
EncoderException("Error encoding %d as signed varint." % value),
exc)
return output
def encode_uvarint(value):
"""Encode a long or int into a bytearray."""
output = bytearray()
if value < 0:
raise EncoderException(
"Error encoding %d as uvarint. Value must be positive" % value)
try:
encoder._EncodeVarint(_gen_append_bytearray(output), value)
except (struct.error, ValueError) as exc:
six.raise_from(
EncoderException("Error encoding %d as uvarint." % value), exc)
return output
def _append_bytearray(x):
if isinstance(x, (str, int)):
arr.append(x)
elif isinstance(x, bytes):
arr.extend(x)
else:
raise EncoderException("Unknown type returned by protobuf library")
def encode_struct(fmt, value):
"""Generic method for encoding arbitrary python "struct" values"""
try:
return struct.pack(fmt, value)
except struct.error as exc:
six.raise_from(
EncoderException("Error encoding value %r with format string %s" %
(value, fmt)), exc)
def encode_string(value):
try:
value = six.ensure_text(value)
except TypeError as exc:
six.raise_from(
EncoderException("Error encoding string to message: %r" % value),
exc)
return encode_bytes(value)
def encode_bytes_hex(value):
"""Encode varint length followed by the string.
Expects a string of hex characters
"""
try:
return encode_bytes(binascii.unhexlify(value))
except (TypeError, binascii.Error) as exc:
six.raise_from(EncoderException("Error encoding hex bytestring %s" % value), exc)
def encode_bytes(value):
"""Encode varint length followed by the string.
This should also work to encode incoming string values.
"""
try:
value = six.ensure_binary(value)
except TypeError as exc:
six.raise_from(EncoderException("Error encoding bytes to message: %r" % value), exc)
encoded_length = varint.encode_varint(len(value))
return encoded_length + value
def _get_field_key(field_number, typedef, path):
"""If field_number has a name, then use that"""
if not isinstance(field_number, (int, str)):
# Should only get unpredictable inputs from encoding
raise EncoderException("Field key in message must be a str or int",
path=path)
if isinstance(field_number, int):
field_number = str(field_number)
alt_field_number = None
if "-" in field_number:
field_number, alt_field_number = field_number.split("-")
# TODO
raise NotImplementedError(
"Handling for _get_field_key not implemented for alt typedefs: %s"
% field_number,
path=path,
)
if field_number in typedef and "name" in typedef[field_number]:
field_key = typedef[field_number]["name"]
else:
field_key = field_number
# Return the new field_name + alt_field_number
return field_key + ("" if alt_field_number is None else "-" +
alt_field_number)
def encode_message(data, config, typedef, path=None):
"""Encode a Python dictionary representing a protobuf message
data - Python dictionary mapping field numbers to values
typedef - Type information including field number, field name and field type
This will throw an exception if an unkown value is used as a key
"""
output = bytearray()
if path is None:
path = []
for field_number, value in data.items():
# Get the field number convert it as necessary
alt_field_number = None
if six.PY2:
string_types = (str, unicode)
else:
string_types = str
if isinstance(field_number, string_types):
if "-" in field_number:
field_number, alt_field_number = field_number.split("-")
# TODO can refactor to cache the name to number mapping
for number, info in typedef.items():
if ("name" in info and info["name"] == field_number
and field_number != ""):
field_number = number
break
else:
field_number = str(field_number)
field_path = path[:]
field_path.append(field_number)
if field_number not in typedef:
raise EncoderException(
"Provided field name/number %s is not valid" % (field_number),
field_path,
)
field_typedef = typedef[field_number]
# Get encoder
if "type" not in field_typedef:
raise TypedefException(
"Field %s does not have a defined type." % field_number,
field_path)
field_type = field_typedef["type"]
field_encoder = None
if alt_field_number is not None:
if alt_field_number not in field_typedef["alt_typedefs"]:
raise EncoderException(
"Provided alt field name/number %s is not valid for field_number %s"
% (alt_field_number, field_number),
field_path,
)
if isinstance(field_typedef["alt_typedefs"][alt_field_number],
dict):
innertypedef = field_typedef["alt_typedefs"][alt_field_number]
field_encoder = lambda data: encode_lendelim_message(
data, config, innertypedef, path=field_path)
else:
# just let the field
field_type = field_typedef["alt_typedefs"][alt_field_number]
if field_encoder is None:
if field_type == "message":
innertypedef = None
if "message_typedef" in field_typedef:
innertypedef = field_typedef["message_typedef"]
elif "message_type_name" in field_typedef:
message_type_name = field_typedef["message_type_name"]
if message_type_name not in config.known_types:
raise TypedefException(
"Message type (%s) has not been defined" %
field_typedef["message_type_name"],
field_path,
)
innertypedef = config.known_types[message_type_name]
else:
raise TypedefException(
"Could not find message typedef for %s" % field_number,
field_path,
)
field_encoder = lambda data: encode_lendelim_message(
data, config, innertypedef, path=field_path)
else:
if field_type not in blackboxprotobuf.lib.types.ENCODERS:
raise TypedefException("Unknown type: %s" % field_type)
field_encoder = blackboxprotobuf.lib.types.ENCODERS[field_type]
if field_encoder is None:
raise TypedefException(
"Encoder not implemented for %s" % field_type,
field_path)
# Encode the tag
tag = encoder.TagBytes(
int(field_number),
blackboxprotobuf.lib.types.WIRETYPES[field_type])
try:
# Handle repeated values
if isinstance(value,
list) and not field_type.startswith("packed_"):
for repeated in value:
output += tag
output += field_encoder(repeated)
else:
output += tag
output += field_encoder(value)
except EncoderException as exc:
exc.set_path(field_path)
six.reraise(*sys.exc_info())
return output
def encode_message(data, typedef, group=False, path=None):
"""Encode a Python dictionary representing a protobuf message
data - Python dictionary mapping field numbers to values
typedef - Type information including field number, field name and field type
This will throw an exception if an unkown value is used as a key
"""
output = bytearray()
if path is None:
path = []
# TODO Implement path for encoding
for field_number, value in data.items():
# Get the field number convert it as necessary
alt_field_number = None
if six.PY2:
string_types = (str, unicode)
else:
string_types = str
if isinstance(field_number, string_types):
if '-' in field_number:
field_number, alt_field_number = field_number.split('-')
for number, info in typedef.items():
if 'name' in info and info['name'] == field_number and field_number != '':
field_number = number
break
else:
field_number = str(field_number)
field_path = path[:]
field_path.append(field_number)
if field_number not in typedef:
raise EncoderException("Provided field name/number %s is not valid"
% (field_number), field_path)
field_typedef = typedef[field_number]
# Get encoder
if 'type' not in field_typedef:
raise TypedefException('Field %s does not have a defined type.' % field_number, field_path)
field_type = field_typedef['type']
field_encoder = None
if field_type == 'message':
innertypedef = None
# Check for a defined message type
if alt_field_number is not None:
if alt_field_number not in field_typedef['alt_typedefs']:
raise EncoderException(
'Provided alt field name/number %s is not valid for field_number %s'
% (alt_field_number, field_number), field_path)
innertypedef = field_typedef['alt_typedefs'][alt_field_number]
elif 'message_typedef' in field_typedef:
innertypedef = field_typedef['message_typedef']
else:
if field_typedef['message_type_name'] not in blackboxprotobuf.lib.known_messages:
raise TypedefException("Message type (%s) has not been defined"
% field_typedef['message_type_name'], field_path)
innertypedef = field_typedef['message_type_name']
field_encoder = lambda data: encode_lendelim_message(data, innertypedef, path=field_path)
elif field_type == 'group':
innertypedef = None
# Check for a defined group type
if 'group_typedef' not in field_typedef:
raise TypedefException("Could not find type definition for group field: %s"
% field_number, field_path)
innertypedef = field_typedef['group_typedef']
field_encoder = lambda data: encode_group(data, innertypedef, field_number, path=field_path)
else:
if field_type not in blackboxprotobuf.lib.types.encoders:
raise TypedefException('Unknown type: %s' % field_type)
field_encoder = blackboxprotobuf.lib.types.encoders[field_type]
if field_encoder is None:
raise TypedefException('Encoder not implemented for %s' % field_type, field_path)
# Encode the tag
tag = encoder.TagBytes(int(field_number), blackboxprotobuf.lib.types.wiretypes[field_type])
try:
# Handle repeated values
if isinstance(value, list) and not field_type.startswith('packed_'):
for repeated in value:
output += tag
output += field_encoder(repeated)
else:
output += tag
output += field_encoder(value)
except EncoderException as exc:
exc.set_path(field_path)
six.reraise(*sys.exc_info())
return output