From 7bb7c2b59bf8ed93a34bee850f71d65e64db54fc Mon Sep 17 00:00:00 2001
From: Simon Grondin <simon.grondin@outlook.com>
Date: Wed, 19 Jul 2023 08:33:06 -0500
Subject: [PATCH] Add Process is_success

---
 lib_eio/process.ml  | 12 ++++++------
 lib_eio/process.mli | 13 +++++++++++--
 tests/process.md    | 22 ++++++++++++++++++++++
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/lib_eio/process.ml b/lib_eio/process.ml
index 405a05f8d..02267efc8 100644
--- a/lib_eio/process.ml
+++ b/lib_eio/process.ml
@@ -39,9 +39,9 @@ end
 let pid proc = proc#pid
 let await proc = proc#await
 
-let await_exn proc =
+let await_exn ?(is_success = Int.equal 0) proc =
   match proc#await with
-  | `Exited 0 -> ()
+  | `Exited code when is_success code -> ()
   | status -> raise (err (Child_error status))
 
 let signal proc = proc#signal
@@ -84,18 +84,18 @@ let spawn ~sw (t:#mgr) ?cwd ?stdin ?stdout ?stderr ?env ?executable args =
     ?stdout:(stdout :> Flow.sink option)
     ?stderr:(stderr :> Flow.sink option)
 
-let run (t:#mgr) ?cwd ?stdin ?stdout ?stderr ?env ?executable args =
+let run (t:#mgr) ?cwd ?stdin ?stdout ?stderr ?(is_success = Int.equal 0) ?env ?executable args =
   Switch.run @@ fun sw ->
   let child = spawn ~sw t ?cwd ?stdin ?stdout ?stderr ?env ?executable args in
   match await child with
-  | `Exited 0 -> ()
+  | `Exited code when is_success code -> ()
   | status ->
     let ex = err (Child_error status) in
     raise (Exn.add_context ex "running command: %a" pp_args args)
 
 let pipe ~sw (t:#mgr) = t#pipe ~sw
 
-let parse_out (t:#mgr) parse ?cwd ?stdin ?stderr ?env ?executable args =
+let parse_out (t:#mgr) parse ?cwd ?stdin ?stderr ?is_success ?env ?executable args =
   Switch.run @@ fun sw ->
   let r, w = pipe t ~sw in
   try
@@ -103,7 +103,7 @@ let parse_out (t:#mgr) parse ?cwd ?stdin ?stderr ?env ?executable args =
     Flow.close w;
     let output = Buf_read.parse_exn parse r ~max_size:max_int in
     Flow.close r;
-    await_exn child;
+    await_exn ?is_success child;
     output
   with Exn.Io _ as ex ->
     let bt = Printexc.get_raw_backtrace () in
diff --git a/lib_eio/process.mli b/lib_eio/process.mli
index 915b8c478..4d6325f77 100644
--- a/lib_eio/process.mli
+++ b/lib_eio/process.mli
@@ -49,8 +49,12 @@ val pid : #t -> int
 val await : #t -> exit_status
 (** [await t] waits for process [t] to exit and then reports the status. *)
 
-val await_exn : #t -> unit
-(** Like {! await} except an exception is raised if the status is not [`Exited 0]. *)
+val await_exn : ?is_success:(int -> bool) -> #t -> unit
+(** Like {! await} except an exception is raised if does not return a successful
+    exit status.
+
+    @param is_success Used to determine if an exit code is successful.
+                      Default is [Int.equal 0]. *)
 
 val signal : #t -> int -> unit
 (** [signal t i] sends the signal [i] to process [t].
@@ -113,12 +117,16 @@ val run :
   ?stdin:#Flow.source ->
   ?stdout:#Flow.sink ->
   ?stderr:#Flow.sink ->
+  ?is_success:(int -> bool) ->
   ?env:string array ->
   ?executable:string ->
   string list -> unit
 (** [run] does {!spawn} followed by {!await_exn}, with the advantage that if the process fails then
     the error message includes the command that failed.
 
+    When [is_success] is provided, it is called with the exit code to determine whether it indicates success or failure.
+    Without [is_success], success requires the process to return an exit code of 0.
+
     Note: If [spawn] needed to create extra fibers to copy [stdin], etc, then it also waits for those to finish. *)
 
 val parse_out :
@@ -127,6 +135,7 @@ val parse_out :
   ?cwd:#Fs.dir Path.t ->
   ?stdin:#Flow.source ->
   ?stderr:#Flow.sink ->
+  ?is_success:(int -> bool) ->
   ?env:string array ->
   ?executable:string ->
   string list -> 'a
diff --git a/tests/process.md b/tests/process.md
index 340a2b4bd..d7427604e 100644
--- a/tests/process.md
+++ b/tests/process.md
@@ -148,6 +148,28 @@ Eio.Io Process Child_error Exited (code 3),
   running command: bash -c "exit 3" "" foo
 ```
 
+Exit code success can be determined by is_success (Process.run):
+
+```ocaml
+# run @@ fun mgr env ->
+  Process.run ~is_success:(Int.equal 3) mgr ["bash"; "-c"; "exit 3"];;
+- : unit = ()
+
+# run @@ fun mgr env ->
+  Process.run ~is_success:(Int.equal 3) mgr ["bash"; "-c"; "exit 0"];;
+Exception:
+Eio.Io Process Child_error Exited (code 0),
+  running command: bash -c "exit 0"
+```
+
+Exit code success can be determined by is_success (Process.parse_out):
+
+```ocaml
+# run @@ fun mgr env ->
+  Process.parse_out ~is_success:(Int.equal 5) mgr Eio.Buf_read.line ["sh"; "-c"; "echo 123; exit 5"];;
+- : string = "123"
+```
+
 The default environment:
 
 ```ocaml