expand(): correct leading zeros issue

Problem: As noted in LLNL#32, the expand() function doesn't correctly
handle hostlists with leading zeros in the ranges.

Rework and restructure the expand() function to correctly handle ranges
of leading zeros, particularly when there are multiple ranges.

Add a unit test using the example provided in LLNL#32 to confirm it is
working as expected.

Author: cmoussa1

hostlist/hostlist.py

@@ -73,67 +73,39 @@ def expand(nodelist):
     :return: The expanded hostlist string.
     """
     node_list = nodelist.split(", ")
-    # print node_list
-
     result_hostlist = []
+
     for node in node_list:
         nodelist_match = r"(\w+-?)\[((,?[0-9]+-?,?-?){0,})\](.*)?"
-        if re.search(nodelist_match, node):
-            match = re.search(nodelist_match, node)
+        match = re.search(nodelist_match, node)
 
-            # holds the ranges of nodes as a string
-            # now we can manipulate the string and cast it to a list of numbers
-            oldstr = str(match.group(2))
-            left_br = oldstr.replace("[", "")
-            right_br = left_br.replace("]", "")
-            num_list = right_br.split(',')
+        if match:
+            prefix = match.group(1)  # hostname prefix
+            ranges_str = match.group(2)  # the range of numbers within brackets
+            suffix = match.group(4) or ""  # optional suffix
 
-            # if the node numbers contain leading zeros, store them to be
-            # prepended in the final list
-            final_list = []
-            lead_zeros = 0
-            lead_zeros_str = ''
-            for elem in num_list:
-                # if it is a range of numbers, break it by the hyphen and
-                # create a list
-                #
-                # will then be merged with final list
-                if '-' in elem:
-                    tmp_list = elem.replace("-", ",").split(",")
-
-                    for digit in tmp_list[0]:
-                        if digit == '0':
-                            lead_zeros = lead_zeros + 1
-                            lead_zeros_str = lead_zeros_str + '0'
-
-                    rng_list = range(int(tmp_list[0]), int(tmp_list[1]) + 1)
-                    final_list.extend(rng_list)
-                else:
-                    final_list.append(int(elem))
-
-            # put final list in ascending order and append cluster name to
-            # each node number
-            final_list.sort()
-
-            # prepend leading zeros to numbers required
-            hostlist_tmp = []
-            for elem in final_list:
-                if ((lead_zeros > 0) and (len(str(elem)) <= len(lead_zeros_str))):
-                    hostlist_tmp.append(str(elem).zfill(lead_zeros + 1))
+            final_hostlist = []
+            ranges = ranges_str.split(',')
+
+            for rng in ranges:
+                if '-' in rng:
+                    start, end = rng.split('-')
+                    # handle leading zeros
+                    width = len(start)
+                    start, end = int(start), int(end)
+                    for i in range(start, end + 1):
+                        final_hostlist.append(f"{prefix}{str(i).zfill(width)}{suffix}")
                 else:
-                    hostlist_tmp.append(str(elem))
+                    final_hostlist.append(f"{prefix}{rng.zfill(len(rng))}{suffix}")
 
-            # append hostname to the node numbers
-            hostlist_no_suffix = []
-            for elem in hostlist_tmp:
-                hostlist_no_suffix.append(match.group(1) + elem)
+            # append sorted hostnames
+            result_hostlist.extend(final_hostlist)
 
-            # append suffix to hostlist if there is one
-            final_hostlist = []
-            for elem in hostlist_no_suffix:
-                final_hostlist.append(elem + match.group(4))
+    # sort the entire hostlist numerically based on the numeric part
+    def numeric_key(hostname):
+        return int(re.search(r'(\d+)', hostname).group(1))
 
-            result_hostlist.append('%s' % ','.join(map(str, final_hostlist)))
+    result_hostlist.sort(key=numeric_key)
 
     return ','.join(result_hostlist)
 

hostlist/unittest_hostlist.py

@@ -20,6 +20,12 @@ def test_expand(self):
         test = hl.expand('quartz[4-8]')
         self.assertEqual(test, expected)
 
+    # expand() with the prompt in issue#32
+    def test_expand_leading_zeros(self):
+        expected = 's02p017,s02p029,s02p030'
+        test = hl.expand('s02p[017,029-030]')
+        self.assertEqual(test, expected)
+
     # expand() will also return correctly with
     # multiple sets of ranges
     def test_expand_multi_range(self):