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