From 5ada2d52cbe6909cbc723c180326172d5db75a4c Mon Sep 17 00:00:00 2001
From: Christopher Moussa <moussa1@llnl.gov>
Date: Wed, 4 Sep 2024 09:59:01 -0700
Subject: [PATCH] expand(): correct leading zeros issue

Problem: As noted in #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 #32 to confirm it is
working as expected.
---
 hostlist/hostlist.py          | 78 +++++++++++------------------------
 hostlist/unittest_hostlist.py |  6 +++
 2 files changed, 30 insertions(+), 54 deletions(-)

diff --git a/hostlist/hostlist.py b/hostlist/hostlist.py
index d58a206..1f59a36 100755
--- a/hostlist/hostlist.py
+++ b/hostlist/hostlist.py
@@ -73,69 +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'
-                        else:
-                            break
-
-                    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)
 
diff --git a/hostlist/unittest_hostlist.py b/hostlist/unittest_hostlist.py
index 78ce542..ee24d22 100755
--- a/hostlist/unittest_hostlist.py
+++ b/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):