diff --git a/hostlist/hostlist.py b/hostlist/hostlist.py index d58a206..22cf038 100755 --- a/hostlist/hostlist.py +++ b/hostlist/hostlist.py @@ -7,13 +7,12 @@ def append_hostname(machine_name, num_list): """ - Helper method to append the hostname to node numbers. + Append the hostname to node numbers. :param machine_name: The name of the cluster. :param num_list: The list of nodes to be appended to the cluster name. :return: A hostlist string with the hostname and node numbers. """ - hostlist = [] for elem in num_list: hostlist.append(machine_name + str(elem)) @@ -23,43 +22,28 @@ def append_hostname(machine_name, num_list): def sort_nodes(nodelist): """ - sort_nodes is a helper method that sorts the nodes in ascending order. + Sort the nodes in a nodelist in ascending order. :param nodelist: The hostlist string. :return: The hostlist string in ascending order. """ + if isinstance(nodelist, str): + # remove brackets and split by commas + nodelist = nodelist.replace("[", "").replace("]", "").split(',') - list_of_nodes = nodelist - result_hostlist = [] - - if type(list_of_nodes) == str: - left_br = list_of_nodes.replace("[", "") - right_br = left_br.replace("]", "") - nodelist = right_br.split(',') - - count = 0 - num_list = [] - for node in nodelist: - iter_node = nodelist[count] - nodelist_match = r"([a-z]+)(\d+)(.*)" - machine_name = re.search(nodelist_match, iter_node) - num_list.append(int(machine_name.group(2))) - count = count + 1 - num_list.sort() - - # append hostname to the node numbers - hostlist_no_suffix = [] - for elem in num_list: - hostlist_no_suffix.append(machine_name.group(1) + str(elem)) - - # append suffix to hostlist if there is one - final_hostlist = [] - for elem in hostlist_no_suffix: - final_hostlist.append(elem + machine_name.group(3)) + def extract_key(node): + nodelist_match = r"([a-zA-Z]+)(\d+)(.*)" + match = re.search(nodelist_match, node) + prefix = match.group(1) + number = int(match.group(2)) + suffix = match.group(3) + return (prefix, number, suffix) - result_hostlist.append('%s' % ','.join(map(str, final_hostlist))) + # sort nodes based on extracted keys (prefix, numeric part, suffix) + sorted_nodes = sorted(nodelist, key=extract_key) - return '%s' % ','.join(map(str, final_hostlist)) + # join the sorted nodes back into a string + return ','.join(sorted_nodes) # ========== END OF HELPER METHODS =========== #