def testEncodeDecodeKey(self):
test_cases = [
'a', 'simple', 'dollar$', '$dollar', '$do$ll$ar$', ('a'),
('a', 'simple'), ('dollar$', 'simple'), ('do$llar', 'sim$ple', 'str$'),
('many', 'many', 'elements', 'in', 'the', 'tuple'), u'unicode\u1234',
u'uni\u1234code\u2345', ('mixed', u'uni\u1234',
u'\u2345\u1234'), (u'\u1234\u2345',
u'\u3456\u2345')
]
for key in test_cases:
self.assertEqual(key, encoding.decode_key(encoding.encode_key(key)))
def get_node_map(
meta_graph_def: meta_graph_pb2.MetaGraphDef, prefix: str,
node_suffixes: List[str]
) -> Dict[types.FPLKeyType, Dict[str, CollectionDefValueType]]:
"""Get node map from meta_graph_def.
This is designed to extract structures of the following form from the
meta_graph_def collection_def:
prefix/key
key1
key2
key3
prefix/suffix_a
node1
node2
node3
prefix/suffix_b
node4
node5
node6
which will become a dictionary:
{
key1 : {suffix_a: node1, suffix_b: node4}
key2 : {suffix_a: node2, suffix_b: node5}
key3 : {suffix_a: node3, suffix_b: node6}
}.
Keys must always be bytes. Values can be any supported CollectionDef type
(bytes_list, any_list, etc)
Args:
meta_graph_def: MetaGraphDef containing the CollectionDefs to extract the
structure from.
prefix: Prefix for the CollectionDef names.
node_suffixes: The suffixes to the prefix to form the names of the
CollectionDefs to extract the nodes from, e.g. in the example described
above, node_suffixes would be ['suffix_a', 'suffix_b'].
Returns:
A dictionary of dictionaries, as described in the example above.
Raises:
ValueError: The length of some node list did not match length of the key
list.
"""
node_lists = []
for node_suffix in node_suffixes:
collection_def_name = encoding.with_suffix(prefix, node_suffix)
collection_def = meta_graph_def.collection_def.get(collection_def_name)
if collection_def is None:
# If we can't find the CollectionDef, append an empty list.
#
# Either all the CollectionDefs are missing, in which case we correctly
# return an empty dict, or some of the CollectionDefs are non-empty,
# in which case we raise an exception below.
node_lists.append([])
else:
node_lists.append(
getattr(collection_def,
collection_def.WhichOneof('kind')).value)
keys = meta_graph_def.collection_def[encoding.with_suffix(
prefix, encoding.KEY_SUFFIX)].bytes_list.value
if not all([len(node_list) == len(keys) for node_list in node_lists]):
raise ValueError(
'length of each node_list should match length of keys. '
'prefix was %s, node_lists were %s, keys was %s' %
(prefix, node_lists, keys))
result = {}
for key, elems in zip(keys, zip(*node_lists)):
result[encoding.decode_key(key)] = dict(zip(node_suffixes, elems))
return result