Stream: add total

Author: arvidj

lib/gitlab_core.ml

@@ -558,6 +558,7 @@ struct
       buffer : 'a list;
       refill : (unit -> 'a t Monad.t) option;
       endpoint : Endpoint.t;
+      total : int option;
     }
 
     type 'a parse = string -> 'a list Lwt.t
@@ -568,6 +569,7 @@ struct
         buffer = [];
         refill = None;
         endpoint = Endpoint.empty;
+        total = Some 0;
       }
 
     let rec next =
@@ -578,6 +580,8 @@ struct
         | { buffer = h :: buffer; _ } as s ->
             return (Some (h, { s with buffer })))
 
+    let total s = s.total
+
     let take : int -> 'a t -> 'b t = fun n s ->
       let rec refill n' s () =
         Monad.(
@@ -645,7 +649,7 @@ struct
       in
       aux [] s
 
-    let of_list buffer = { empty with buffer; refill = None }
+    let of_list buffer = { empty with buffer; refill = None; total = Some (List.length buffer) }
     let poll stream = stream.restart stream.endpoint
 
     let since stream version =
@@ -848,6 +852,7 @@ struct
                 buffer;
                 refill = None;
                 endpoint = Endpoint.empty;
+                total = None;
               })
         fhs
 
@@ -867,8 +872,13 @@ struct
               | Some uri ->
                   request ~uri (stream_next restart request uri fn endpoint))
         in
+        let total =
+          Option.bind
+            Cohttp.(Header.get envelope.Response.headers "x-total")
+            int_of_string_opt
+        in
         fn body >>= fun buffer ->
-        return { Stream.restart; buffer; refill; endpoint })
+        return { Stream.restart; buffer; refill; endpoint; total })
 
     let rec restart_stream ~fail_handlers ~expected_code ?headers ?token ?params
         fn endpoint =
@@ -933,7 +943,7 @@ struct
         Some
           (fun () -> request ~uri (stream_next restart request uri fn endpoint))
       in
-      { Stream.restart; buffer = []; refill; endpoint }
+      { Stream.restart; buffer = []; refill; endpoint; total = None }
 
     let post ?(fail_handlers = []) ~expected_code ?headers =
       let headers =

lib/gitlab_s.mli

@@ -212,6 +212,8 @@ module type Gitlab = sig
         means to iterate over ordered API results which may be too
         numerous to fit into a single request/response pair. *)
 
+    val total : 'a t -> int option
+
     val map : ('a -> 'b list Monad.t) -> 'a t -> 'b t
     (** [map f s] is the lazy stream of [f] applied to elements of [s]
         as they are demanded. *)