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