keep_header source code

1 /**
2 Command line tool that executes a command while preserving header lines.
3 
4 Copyright (c) 2018, eBay Software Foundation
5 Initially written by Jon Degenhardt
6 
7 License: Boost License 1.0 (http://boost.org/LICENSE_1_0.txt)
8 */
9 module 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 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 int main(string[] args)
36 {
37     import std.algorithm : findSplit, joiner;
38     import std.path : baseName, stripExtension;
39     import std.process : pipeProcess, ProcessPipes, Redirect, wait;
40     import std.range;
41     import std.stdio;
42     import std.typecons : tuple;
43 
44     /* When running in DMD code coverage mode, turn on report merging. */
45     version(D_Coverage) version(DigitalMars)
46     {
47         import core.runtime : dmd_coverSetMerge;
48         dmd_coverSetMerge(true);
49     }
50 
51     auto programName = (args.length > 0) ? args[0].stripExtension.baseName : "Unknown_program_name";
52     auto splitArgs = findSplit(args, ["--"]);
53 
54     if (splitArgs[1].length == 0 || splitArgs[2].length == 0)
55     {
56         auto cmdArgs = splitArgs[0][1 .. $];
57         stderr.writefln("Synopsis: %s [file...] -- program [args]", programName);
58         if (cmdArgs.length > 0 &&
59             (cmdArgs[0] == "-h" || cmdArgs[0] == "--help" || cmdArgs[0] == "--help-verbose"))
60         {
61             stderr.writeln();
62             stderr.writeln(helpText);
63         }
64         else if (cmdArgs.length > 0 &&
65                  (cmdArgs[0] == "-V" || cmdArgs[0] == "--V" ||  cmdArgs[0] == "--version"))
66         {
67             import tsvutils_version;
68             stderr.writeln();
69             stderr.writeln(tsvutilsVersionNotice("keep-header"));
70         }
71         return 0;
72     }
73 
74     ProcessPipes pipe;
75     try pipe = pipeProcess(splitArgs[2], Redirect.stdin);
76     catch (Exception exc)
77     {
78         stderr.writefln("[%s] Command failed: '%s'", programName, splitArgs[2].joiner(" "));
79         stderr.writeln(exc.msg);
80         return 1;
81     }
82 
83     int status = 0;
84     {
85         scope(exit)
86         {
87             auto pipeStatus = wait(pipe.pid);
88             if (pipeStatus != 0) status = pipeStatus;
89         }
90 
91         bool headerWritten = false;
92         foreach (filename; splitArgs[0].length > 1 ? splitArgs[0][1..$] : ["-"])
93         {
94             bool isStdin = (filename == "-");
95             File inputStream;
96 
97             if (isStdin) inputStream = stdin;
98             else
99             {
100                 try inputStream = filename.File();
101                 catch (Exception exc)
102                 {
103                     stderr.writefln("[%s] Unable to open file: '%s'", programName, filename);
104                     stderr.writeln(exc.msg);
105                     status = 1;
106                     break;
107                 }
108             }
109 
110             auto firstLine = inputStream.readln();
111 
112             if (inputStream.eof && firstLine.length == 0) continue;
113 
114             if (!headerWritten)
115             {
116                 write(firstLine);
117                 stdout.flush;
118                 headerWritten = true;
119             }
120 
121             if (isStdin)
122             {
123                 foreach (line; inputStream.byLine(KeepTerminator.yes))
124                 {
125                     pipe.stdin.write(line);
126                 }
127             }
128             else
129             {
130                 ubyte[1024*1024] readBuffer;
131                 foreach (ubyte[] chunk; inputStream.byChunk(readBuffer))
132                 {
133                     pipe.stdin.write(cast(char[])chunk);
134                 }
135             }
136             pipe.stdin.flush;
137         }
138         pipe.stdin.close;
139     }
140     return status;
141 }