def _GetBaseApiMap(self):
"""Get a map from graph op name to its base ApiDef.
Returns:
Dictionary mapping graph op name to corresponding ApiDef.
"""
# Convert base ApiDef in Multiline format to Proto format.
converted_base_api_dir = os.path.join(
test.get_temp_dir(), 'temp_base_api_defs')
subprocess.check_call(
[os.path.join(resource_loader.get_root_dir_with_all_resources(),
_CONVERT_FROM_MULTILINE_SCRIPT),
_BASE_API_DIR, converted_base_api_dir])
name_to_base_api_def = {}
base_api_files = file_io.get_matching_files(
os.path.join(converted_base_api_dir, 'api_def_*.pbtxt'))
for base_api_file in base_api_files:
if file_io.file_exists(base_api_file):
api_defs = api_def_pb2.ApiDefs()
text_format.Merge(
file_io.read_file_to_string(base_api_file), api_defs)
for api_def in api_defs.op:
name_to_base_api_def[api_def.graph_op_name] = api_def
return name_to_base_api_def
def read_api_def_map(folder):
api_def_map = {}
file_list = os.listdir(folder)
file_list = sorted(file_list)
for filename in file_list:
api_defs = api_def_pb2.ApiDefs()
filename = folder + '/' + filename
with open(filename) as handle:
multiline_pbtxt = handle.read()
pbtxt = pbtxt_from_multiline(multiline_pbtxt)
text_format.Merge(pbtxt, api_defs)
for api_def in api_defs.op:
api_def_map[api_def.graph_op_name] = api_def
return api_def_map
