Input from stdin and argument in Go

» Introduction

When I wrote my last two tools multicode and epoch, I wanted them to work with input arguments, as well as used in a chain of a pipeline (a Unix pipe). Here are examples of the two input methods.

Let’s assume we have the input Njg2NTZjNmM2ZjIwNzQ2ODY1NzI2NTBhCg== for the command decode.

Using the argument, the input data is just added after the command name:

1
2

$ decode Njg2NTZjNmM2ZjIwNzQ2ODY1NzI2NTBhCg==
hello there

Using the Unix pipe |, the input data is the result of the previous command. Here, the example is very synthetic, as the first command (echo), is just returning the given input unmodified. The result of echo, will be the input of the next command (decode).

1
2

$ echo Njg2NTZjNmM2ZjIwNzQ2ODY1NzI2NTBhCg== | decode
hello there

Here is a second example using the pipe command. This time with four commands.

1
2

$ echo hello there | xxd -p | base64 | decode
hello there

And another representation of the above line with the intermediate results for a better understanding:

1. echo hello there -> hello there
2. hello there -> xxd -p -> 68656c6c6f2074686572650a
3. 68656c6c6f2074686572650a -> base64 -> Njg2NTZjNmM2ZjIwNzQ2ODY1NzI2NTBhCg==
4. Njg2NTZjNmM2ZjIwNzQ2ODY1NzI2NTBhCg== -> decode -> hello there

» Analysing gofmt

How can we make our program use both, input from arguments and the pipe? After I’ve searched for it and did not find any results, I thought that I already know at least one Go tool which works like this: gofmt (which formats go code to a specific style). You can use it and give a file name as an argument, or input the code directly, using the pipe. So the functionality is a bit different between the two options. Reading a file and processing the content, or processing the content without the step of reading a file first.

gofmt will read the given file(s):

1

gofmt main.go

gofmt uses the source code directly as printed by cat.

1

cat main.go | gofmt

So, I’ve just checked the source of gofmt and found the right place in the code:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

if flag.NArg() == 0 {
    ...
    if err := processFile("<standard input>", os.Stdin, os.Stdout, true); err != nil {
        report(err)
    }
    return
}

for i := 0; i < flag.NArg(); i++ {
    path := flag.Arg(i)
    switch dir, err := os.Stat(path); {
    case err != nil:
        report(err)
    case dir.IsDir():
        walkDir(path)
    default:
        if err := processFile(path, nil, os.Stdout, false); err != nil {
            report(err)
        }
    }
}

Basically, all of the magic is checking the number of input arguments (flag.NArg()). When there are no arguments, use the values from stdin (the data from the pipe or directly input. If directly input, use CTRL-D to signal the end of the input).

If the number of arguments is not zero, check if the input is a file (or directory) and process with the data from the given path instead of using stdin.

You can see, how the processFile function allows handling both inputs. Here is the signature of the function:

1

func processFile(filename string, in io.Reader, out io.Writer, stdin bool) error

So, it contains the logic of first reading a file, or skipping this step and using the input directly.

» Minimal Example

Here is how I adopted the code for my tools. I have an even simpler use case, as I don’t have to read the data from a file, just using the plain input in both cases:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

flag.Parse()
var input string

// read program input
if flag.NArg() == 0 { // from stdin/pipe
    reader := bufio.NewReader(os.Stdin)
    var err error
    input, err = reader.ReadString('\n')
    if err != nil {
        log.Fatalln("failed to read input")
    }
    input = strings.TrimSpace(input) // otherwise, we would have a blank line
} else { // from argument
    if flag.NArg() > 1 {
        log.Fatalln("takes at most one input")
    }
    input = flag.Arg(0)
}

fmt.Printf("> %s\n", input)

After processing the if...else statement, input will contain the same data, regardless which input method was used.

You can find the minimal example program in my playground repo.