Source file `write.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
open Core
open Async
include Delimited_kernel.Write

module Raw = struct
  let of_writer ~init ~f writer =
    Pipe.create_writer (fun reader ->
      init writer;
      Writer.transfer writer reader (f ~writer))
  ;;

  let of_writer_and_close ~init ~f writer =
    let pipe = of_writer ~init ~f writer in
    don't_wait_for
      (let%bind () = Pipe.closed pipe in
       Writer.close writer);
    pipe
  ;;

  let create_writer filename ~init ~f =
    let%map writer = Writer.open_file filename in
    of_writer_and_close writer ~init ~f
  ;;
end

module Expert = struct
  include Delimited_kernel.Write.Expert

  module By_row = struct
    let write_field ~sep w field = Writer.write w (maybe_escape_field ~sep field)

    let write_line ?(sep = ',') ?(line_breaks = `Windows) ~writer line =
      let line_breaks =
        match line_breaks with
        | `Unix -> "\n"
        | `Windows -> "\r\n"
      in
      let rec loop line =
        match line with
        | [] -> Writer.write writer line_breaks
        | [ field ] ->
          write_field ~sep writer field;
          loop []
        | field :: rest ->
          write_field ~sep writer field;
          Writer.write_char writer sep;
          loop rest
      in
      loop line
    ;;

    let base ?sep ?line_breaks create =
      create ~init:(Fn.const ()) ~f:(write_line ?sep ?line_breaks)
    ;;

    let of_writer_and_close ?sep ?line_breaks writer =
      base ?sep ?line_breaks (Raw.of_writer_and_close writer)
    ;;

    let of_writer ?sep ?line_breaks writer = base ?sep ?line_breaks (Raw.of_writer writer)

    let create_writer ?sep ?line_breaks filename =
      base ?sep ?line_breaks (Raw.create_writer filename)
    ;;
  end

  let base ?sep ?line_breaks ~builder ~write_header create =
    let init =
      if write_header
      then fun writer -> By_row.write_line ?sep ?line_breaks ~writer (headers builder)
      else fun (_ : Writer.t) -> ()
    in
    let f ~writer line =
      By_row.write_line ?sep ?line_breaks ~writer (to_columns builder line)
    in
    create ~init ~f
  ;;

  let of_writer ?sep ?line_breaks ~write_header builder writer =
    base ?sep ?line_breaks ~write_header ~builder (Raw.of_writer writer)
  ;;

  let of_writer_and_close ?sep ?line_breaks ~write_header builder writer =
    base ?sep ?line_breaks ~write_header ~builder (Raw.of_writer_and_close writer)
  ;;

  let create_writer ?sep ?line_breaks ~write_header builder filename =
    base ?sep ?line_breaks ~write_header ~builder (Raw.create_writer filename)
  ;;
end

let protect ~f pipe =
  Monitor.protect
    ~run:`Schedule
    ~rest:`Log
    (fun () -> f pipe)
    ~finally:(fun () ->
      Pipe.close pipe;
      Deferred.ignore_m (Pipe.upstream_flushed pipe))
;;

module By_row = struct
  include Delimited_kernel.Write.By_row

  let with_writer ?sep ?line_breaks writer ~f =
    let pipe = Expert.By_row.base ?sep ?line_breaks (Raw.of_writer writer) in
    protect pipe ~f
  ;;

  let with_file ?sep ?line_breaks filename ~f =
    Writer.with_file filename ~f:(fun writer -> with_writer ?sep ?line_breaks writer ~f)
  ;;

  let with_file_atomic ?temp_file ?fsync ?sep ?line_breaks filename ~f =
    Writer.with_file_atomic ?temp_file ?fsync filename ~f:(fun writer ->
      with_writer ?sep ?line_breaks writer ~f)
  ;;
end

let with_writer ?sep ?line_breaks ~write_header builder writer ~f =
  let pipe =
    Expert.base ?sep ?line_breaks ~write_header ~builder (Raw.of_writer writer)
  in
  protect ~f pipe
;;

let with_file ?sep ?line_breaks ~write_header builder filename ~f =
  Writer.with_file filename ~f:(fun writer ->
    with_writer ?sep ?line_breaks ~write_header builder writer ~f)
;;