Something went wrong. Try again.
Eio flows with progress bars
Something went wrong. Try again.
4.5 kB · 137 lines
OCaml
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138open Eio.Std
type 'a report = 'a -> unittype display = (int report, int report) Progress.Display.t
let display_key : display Fiber.key = Fiber.create_key ()let get_display () = Fiber.get display_key
let get_display = function | Some display -> display | None -> ( match get_display () with | Some display -> display | None -> Fmt.failwith "No progress display provided or stored in fiber. Try \ Eio_progress.with_display %@%@ fun display ->.")
let init ?config sw = let display = Progress.Display.start ?config Progress.Multi.blank in Switch.on_release sw (fun () -> Progress.Display.finalise display); display
let with_display ?config fn = Switch.run ~name:"progress-display" @@ fun sw -> let display = init ?config sw in Fiber.with_binding display_key display @@ fun () -> fn display
let of_source ?message ?size ?report ?display (Eio.Resource.T (v, bindings) : 'a Eio.Flow.source) = let module X = (val Eio.Resource.get bindings Eio.Flow.Pi.Source) in let display = get_display display in let report = match (report, size) with | Some r, _ -> r | None, Some size -> let report = Progress.Display.add_line display (Progress.counter ?message size) in fun i -> Progress.Reporter.report report (Int64.of_int i) | None, None -> Fmt.failwith "No size or reporter for source." in let module V = struct include X
let single_read v bufs = let r = X.single_read v bufs in report r; r end in let bindings = Eio.Resource.bindings bindings in let bindings = Eio.Resource.handler (Eio.Resource.H (Eio.Flow.Pi.Source, (module V)) :: bindings) in Eio.Resource.T (v, bindings)
let of_sink report (Eio.Resource.T (v, bindings) : _ Eio.Flow.sink) = let module X = (val Eio.Resource.get bindings Eio.Flow.Pi.Sink) in let module V = struct include X
let single_write v bufs = let r = X.single_write v bufs in report r; r end in let bindings = Eio.Resource.bindings bindings in let bindings = Eio.Resource.handler (Eio.Resource.H (Eio.Flow.Pi.Sink, (module V)) :: bindings) in Eio.Resource.T (v, bindings)
module Path = struct let with_open_in ?bar ?display path fn = Eio.Switch.run ~name:"with_open_in" @@ fun sw -> let display = get_display display in let file_ro = Eio.Path.open_in ~sw path in let size = Eio.File.size file_ro in let line = match bar with | Some bar -> bar | None -> Progress.counter ?message:(Option.map snd @@ Eio.Path.split path) (Optint.Int63.to_int64 size) in let r = Progress.Display.add_line display line in Eio.Switch.on_release sw (fun () -> Progress.Reporter.finalise r); let r = Progress.Reporter.report r in let r i = r (Int64.of_int i) in let flow = of_source ~report:r file_ro in fn ~size flow
let with_open_out ?bar ?display ?append ~size ~create path fn = Eio.Switch.run ~name:"with_open_out" @@ fun sw -> let display = get_display display in let file_wr = Eio.Path.open_out ?append ~create ~sw path in let line = match bar with | Some bar -> bar | None -> Progress.counter ?message:(Option.map snd @@ Eio.Path.split path) (Optint.Int63.to_int64 size) in let r = Progress.Display.add_line display line in Eio.Switch.on_release sw (fun () -> Progress.Reporter.finalise r); let r = Progress.Reporter.report r in let r i = r (Int64.of_int i) in let flow = of_sink r file_wr in fn flow
let load ?bar ?display path = with_open_in ?bar ?display path @@ fun ~size flow -> try if Optint.Int63.(compare size (of_int Sys.max_string_length)) = 1 then raise @@ Eio.Fs.err File_too_large; let buf = Cstruct.create (Optint.Int63.to_int size) in let rec loop buf got = match Eio.Flow.single_read flow buf with | n -> loop (Cstruct.shift buf n) (n + got) | exception End_of_file -> got in let got = loop buf 0 in Cstruct.to_string ~len:got buf with Eio.Exn.Io _ as ex -> let bt = Printexc.get_raw_backtrace () in Eio.Exn.reraise_with_context ex bt "loading %a" Eio.Path.pp path
let save ?bar ?display ?append ~create path s = with_open_out ?bar ?display ?append ~create ~size:(Optint.Int63.of_int (String.length s)) path @@ fun flow -> Eio.Flow.copy_string s flowend