tsv_utils.keep_header source code

1 /**
2 Command line tool that executes a command while preserving header lines.
3 
4 Copyright (c) 2018-2021, eBay Inc.
5 Initially written by Jon Degenhardt
6 
7 License: Boost License 1.0 (http://boost.org/LICENSE_1_0.txt)
8 */
9 module tsv_utils.keep_header;
10 
11 auto helpText = q"EOS
12 Execute a command against one or more files in a header aware fashion.
13 The first line of each file is assumed to be a header. The first header
14 is output unchanged. Remaining lines are sent to the given command via
15 standard input, excluding the header lines of subsequent files. Output
16 from the command is appended to the initial header line.
17 
18 A double dash (--) delimits the command, similar to how the pipe
19 operator (|) delimits commands. Examples:
20 
21     $ keep-header file1.txt -- sort
22     $ keep-header file1.txt file2.txt -- sort -k1,1nr
23 
24 These sort the files as usual, but preserve the header as the first line
25 output. Data can also be read from standard input. Example:
26 
27     $ cat file1.txt | keep-header -- grep red
28 
29 Options:
30 
31 -V      --version   Print version information and exit.
32 -h         --help   This help information.
33 EOS";
34 
35 static if (__VERSION__ >= 2085) extern(C) __gshared string[] rt_options = [ "gcopt=cleanup:none" ];
36 
37 /** keep-header is a simple program, it is implemented entirely in main.
38  */
39 int main(string[] args)
40 {
41     import std.algorithm : findSplit, joiner;
42     import std.path : baseName, stripExtension;
43     import std.process : pipeProcess, ProcessPipes, Redirect, wait;
44     import std.range;
45     import std.stdio;
46     import std.typecons : tuple;
47 
48     /* When running in DMD code coverage mode, turn on report merging. */
49     version(D_Coverage) version(DigitalMars)
50     {
51         import core.runtime : dmd_coverSetMerge;
52         dmd_coverSetMerge(true);
53     }
54 
55     auto programName = (args.length > 0) ? args[0].stripExtension.baseName : "Unknown_program_name";
56     auto splitArgs = findSplit(args, ["--"]);
57 
58     if (splitArgs[1].length == 0 || splitArgs[2].length == 0)
59     {
60         auto cmdArgs = splitArgs[0][1 .. $];
61         stderr.writefln("Synopsis: %s [file...] -- program [args]", programName);
62         if (cmdArgs.length > 0 &&
63             (cmdArgs[0] == "-h" || cmdArgs[0] == "--help" || cmdArgs[0] == "--help-verbose"))
64         {
65             stderr.writeln();
66             stderr.writeln(helpText);
67         }
68         else if (cmdArgs.length > 0 &&
69                  (cmdArgs[0] == "-V" || cmdArgs[0] == "--V" ||  cmdArgs[0] == "--version"))
70         {
71             import tsv_utils.common.tsvutils_version;
72             stderr.writeln();
73             stderr.writeln(tsvutilsVersionNotice("keep-header"));
74         }
75         return 0;
76     }
77 
78     ProcessPipes pipe;
79     try pipe = pipeProcess(splitArgs[2], Redirect.stdin);
80     catch (Exception exc)
81     {
82         stderr.writefln("[%s] Command failed: '%s'", programName, splitArgs[2].joiner(" "));
83         stderr.writeln(exc.msg);
84         return 1;
85     }
86 
87     int status = 0;
88     {
89         scope(exit)
90         {
91             auto pipeStatus = wait(pipe.pid);
92             if (pipeStatus != 0) status = pipeStatus;
93         }
94 
95         bool headerWritten = false;
96         foreach (filename; splitArgs[0].length > 1 ? splitArgs[0][1..$] : ["-"])
97         {
98             bool isStdin = (filename == "-");
99             File inputStream;
100 
101             if (isStdin) inputStream = stdin;
102             else
103             {
104                 try inputStream = filename.File();
105                 catch (Exception exc)
106                 {
107                     stderr.writefln("[%s] Unable to open file: '%s'", programName, filename);
108                     stderr.writeln(exc.msg);
109                     status = 1;
110                     break;
111                 }
112             }
113 
114             auto firstLine = inputStream.readln();
115 
116             if (inputStream.eof && firstLine.length == 0) continue;
117 
118             if (!headerWritten)
119             {
120                 write(firstLine);
121                 stdout.flush;
122                 headerWritten = true;
123             }
124 
125             if (isStdin)
126             {
127                 foreach (line; inputStream.byLine(KeepTerminator.yes))
128                 {
129                     pipe.stdin.write(line);
130                 }
131             }
132             else
133             {
134                 ubyte[1024 * 128] readBuffer;
135                 foreach (ubyte[] chunk; inputStream.byChunk(readBuffer))
136                 {
137                     pipe.stdin.write(cast(char[])chunk);
138                 }
139             }
140             pipe.stdin.flush;
141         }
142         pipe.stdin.close;
143     }
144     return status;
145 }