From 5b04fc26e362cb0fffcf3b1bd94eee0ac5a8aba7 Mon Sep 17 00:00:00 2001
From: Christopher Moussa <moussa1@llnl.gov>
Date: Wed, 4 Sep 2024 11:17:03 -0700
Subject: [PATCH] helper functions: improve helper functions

Problem: The function descriptions for append_hostname() and
sort_nodes() could use some better descriptions. sort_nodes() as a
whole could probably use some improvement via condensing it down.

Edit the descriptions for both helper functions and remove the extra
whitespace between the function description and the start of the
function definition.

In sort_nodes():

- simplify the string processing by checking "nodelist" directly
instead of assigning it to another variable.
- add an extract_key() function to extract the prefix, numeric part,
and suffix using regex and return a tuple of it.
- use the returned tuple from extract_key() to sort the original
nodelist by prefix, numeric part, and suffix.
---
 hostlist/hostlist.py | 48 +++++++++++++++-----------------------------
 1 file changed, 16 insertions(+), 32 deletions(-)

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 =========== #