flux-hostlist: allow idset argument to --nth and --exclude options

doc/man1/flux-hostlist.rst

@@ -63,10 +63,12 @@ OPTIONS
   Emit the number of hosts in the result hostlist instead of the hostlist
   itself.
 
-.. option:: -n, --nth=N
+.. option:: -n, --nth=IDS
 
-  Output only the host at index *N* (*-N* to index from the end). The command
-  will fail if *N* is not a valid index.
+  Output only the hosts at indices *IDS* (*-IDS* to index from the end),
+  where *IDS* is a valid RFC 22 idset (e.g. '0' will return the first host,
+  '0-1' will return the first and second, '-1' returns the last host).  The
+  command will fail if any id in *IDS* is not a valid index.
 
 .. option:: -L, --limit=N
 
@@ -82,9 +84,11 @@ OPTIONS
   other manipulation options, this is equivalent to returning the set
   union of all provided hosts. (By default, all inputs are concatenated).
 
-.. option:: -x, --exclude=HOSTS
+.. option:: -x, --exclude=HOSTS|IDS
 
-  Exclude all occurrences of *HOSTS* form the result.
+  Exclude all hosts in  *HOSTS* or indices in idset *IDS* from the result.
+  It is not an error if any hosts or indices do not exist in the target
+  hostlist.
 
 .. option:: -i, --intersect
 
@@ -148,6 +152,12 @@ List the hosts for one job: (Note: this is the same as
   $ flux hostlist JOBID
   host[1-2]
 
+List the hosts for one job, excluding the first node:
+
+::
+
+  $ flux hostlist -x 0 JOBID
+
 List the unordered, unique hosts for multiple jobs:
 
 ::
@@ -216,6 +226,7 @@ FLUX RFC
 ========
 
 :doc:`rfc:spec_29`
+:doc:`rfc:spec_22`
 
 
 SEE ALSO

src/cmd/flux-hostlist.py

@@ -17,6 +17,7 @@
 import flux
 import flux.util
 from flux.hostlist import Hostlist
+from flux.idset import IDset
 from flux.job import JobID, job_list_id
 from flux.resource import resource_status
 
@@ -64,12 +65,12 @@
         action="store_true",
         help="Print the total number of hosts",
     )
-    group.add_argument(
+    parser.add_argument(
         "-n",
         "--nth",
-        type=int,
-        metavar="N",
-        help="Output host at index N (-N to index from end)",
+        type=str,
+        metavar="[-]IDS",
+        help="Output hosts at indices in idset IDS (-IDS to index from end)",
     )
     parser.add_argument(
         "-L",
@@ -87,9 +88,9 @@
     parser.add_argument(
         "-x",
         "--exclude",
-        metavar="HOSTS",
+        metavar="IDS|HOSTS",
         type=Hostlist,
-        help="Exclude all occurrences of HOSTS from final result",
+        help="Exclude all occurrences of HOSTS or indices from final result",
     )
     parser.add_argument(
         "-u",
@@ -356,8 +357,17 @@
 
     if args.exclude:
         # Delete all occurrences of args.exclude
+        count = len(hl)
         while hl.delete(args.exclude) > 0:
             pass
+        if len(hl) == count:
+            # No hosts were deleted, try args.exclude as idset of indices:
+            try:
+                exclude = IDset(args.exclude)
+                hl = Hostlist([hl[i] for i in range(count) if i not in exclude])
+            except ValueError:
+                # not a valid idset, just pass unaltered hostlist along
+                pass
 
     if args.sort:
         hl.sort()
@@ -375,10 +385,14 @@
         sys.stdout = open(os.devnull, "w")
 
     if args.nth is not None:
-        host = hl[args.nth]
-        if host:
-            print(host)
-    elif args.count:
+        if args.nth.startswith("-"):
+            # Iterate idset in reverse so that resultant hostlist is in
+            # the same order as the input hostlist instead of reversed:
+            hl = Hostlist([hl[-x] for x in reversed(list(IDset(args.nth[1:])))])
+        else:
+            hl = hl[IDset(args.nth)]
+
+    if args.count:
         print(f"{hl.count()}")
     elif args.expand:
         # Convert '\n' specified on command line to actual newline char

t/t2814-hostlist-cmd.t

@@ -100,8 +100,29 @@ test_expect_success 'flux-hostlist -n, --nth works' '
 	test "$(flux hostlist --nth=1 foo[1-10])" = foo2 &&
 	test "$(flux hostlist --nth=-1 foo[1-10])" = foo10
 '
+test_expect_success 'flux-hostlist -n, --nth works with an idset' '
+	flux hostlist --nth=1,3-4 foo[1-10] &&
+	test "$(flux hostlist --nth=1,3-4 foo[1-10])" = "foo[2,4-5]" &&
+	flux hostlist --nth=-1,3-4 foo[1-10] &&
+	test "$(flux hostlist --nth=-1,3-4 foo[1-10])" = "foo[7-8,10]"
+'
+test_expect_success 'flux-hostlist -n, --nth works with --expand' '
+	test "$(flux hostlist -e --nth=1,3-4 foo[1-10])" = "foo2 foo4 foo5"
+'
 test_expect_success 'flux-hostlist -n errors with invalid index' '
-	test_must_fail flux hostlist -n 10 foo[1-10]
+	test_must_fail flux hostlist -n 10 foo[1-10] &&
+	test_must_fail flux hostlist -n 1,10 foo[1-10]
+'
+test_expect_success 'flux-hostlist -x, --exclude works' '
+	test "$(flux hostlist -x foo1 foo[0-10])" = "foo[0,2-10]" &&
+	test "$(flux hostlist -x foo[0-9] foo[0-10])" = "foo10"
+'
+test_expect_success 'flux-hostlist -x, --exclude works with indices' '
+	test "$(flux hostlist -x 1 foo[0-10])" = "foo[0,2-10]" &&
+	test "$(flux hostlist -x 0-9 foo[0-10])" = "foo10"
+'
+test_expect_success 'flux-hostlist -n works after -x' '
+	test "$(flux hostlist -x foo5 -n 5-6 foo[1-10])" = "foo[7-8]"
 '
 test_expect_success 'flux-hostlist -L, --limit works' '
 	test "$(flux hostlist -L 2 foo[1-10])" = "foo[1-2]" &&