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
open Core_kernel
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
    (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)
  ;;
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)
;;