def get_mock_config(config_dict=None):
"""Returns a protobuf Config object from topology_pb2"""
if config_dict is None:
return topology_pb2.Config()
proto_config = topology_pb2.Config()
config_serializer = PythonSerializer()
assert isinstance(config_dict, dict)
for key, value in config_dict.items():
if isinstance(value, bool):
kvs = proto_config.kvs.add()
kvs.key = key
kvs.value = "true" if value else "false"
kvs.type = topology_pb2.ConfigValueType.Value("STRING_VALUE")
elif isinstance(value, (str, int, float)):
kvs = proto_config.kvs.add()
kvs.key = key
kvs.value = str(value)
kvs.type = topology_pb2.ConfigValueType.Value("STRING_VALUE")
else:
kvs = proto_config.kvs.add()
kvs.key = key
kvs.serialized_value = config_serializer.serialize(value)
kvs.type = topology_pb2.ConfigValueType.Value(
"PYTHON_SERIALIZED_VALUE")
return proto_config
def get_mock_config(config_dict=None):
"""Returns a protobuf Config object from topology_pb2"""
if config_dict is None:
return topology_pb2.Config()
proto_config = topology_pb2.Config()
config_serializer = PythonSerializer()
assert isinstance(config_dict, dict)
for key, value in config_dict.items():
if isinstance(value, bool):
kvs = proto_config.kvs.add()
kvs.key = key
kvs.value = "true" if value else "false"
kvs.type = topology_pb2.ConfigValueType.Value("STRING_VALUE")
elif isinstance(value, (str, int, float)):
kvs = proto_config.kvs.add()
kvs.key = key
kvs.value = str(value)
kvs.type = topology_pb2.ConfigValueType.Value("STRING_VALUE")
else:
kvs = proto_config.kvs.add()
kvs.key = key
kvs.serialized_value = config_serializer.serialize(value)
kvs.type = topology_pb2.ConfigValueType.Value("PYTHON_SERIALIZED_VALUE")
return proto_config
def test_py_serializer(self):
serializer = PythonSerializer()
serializer.initialize()
# Test with a list of primitive types
for obj in prim_list:
serialized = serializer.serialize(obj)
self.assertIsInstance(serialized, six.binary_type)
deserialized = serializer.deserialize(serialized)
self.assertEqual(deserialized, obj)
def make_data_tuple_from_list(lst, serializer=PythonSerializer()):
"""Make HeronDataTuple from a list of objects"""
data_tuple = tuple_pb2.HeronDataTuple()
data_tuple.key = 0
tuple_size_in_bytes = 0
for obj in lst:
serialized = serializer.serialize(obj)
data_tuple.values.append(serialized)
tuple_size_in_bytes += len(serialized)
return data_tuple, tuple_size_in_bytes
def test_py_serializer(self):
serializer = PythonSerializer()
serializer.initialize()
# Test with a list of primitive types
for obj in prim_list:
serialized = serializer.serialize(obj)
self.assertIsInstance(serialized, str)
deserialized = serializer.deserialize(serialized)
self.assertEqual(deserialized, obj)
def get_serializer(context):
"""Returns a serializer for a given context"""
cluster_config = context.get_cluster_config()
serializer_clsname = cluster_config.get(constants.TOPOLOGY_SERIALIZER_CLASSNAME, None)
if serializer_clsname is None:
return PythonSerializer()
else:
try:
topo_pex_path = context.get_topology_pex_path()
pex_loader.load_pex(topo_pex_path)
serializer_cls = pex_loader.import_and_get_class(topo_pex_path, serializer_clsname)
serializer = serializer_cls()
return serializer
except Exception as e:
raise RuntimeError("Error with loading custom serializer class: %s, with error message: %s"
% (serializer_clsname, str(e)))
