Golangの標準入力をインタラクティブかパイプで入力を受け取る

Go言語でコマンドラインツールを作るときに入力を受け取るインターフェースでオプションや標準入力で受け付けることはあると思いますが、パイプで渡すことも考慮されているとクールなツールになるなと思っています。

標準入力の受け取り

それぞれの実装方法は簡単です。

インタラクティブ

var stdin string
fmt.Scan(&stdin)
fmt.Println(stdin)

インタラクティブに標準入力からデータを受け取るには fmt.Scan で入力待ちをします。このとき入力した値が渡した変数に格納されます。

パイプ

body, err := ioutil.ReadAll(os.Stdin)
fmt.Println(string(body))

パイプで渡ってきたものは os.Stdin というファイルディスクリプタにデータが入っているので、ここから取得します。

インタラクティブかパイプを判定する

パイプでファイルディスクリプタが渡ってきた場合はそのままそのファイルディスクリプタからデータを取得すれば良いので、インタラクティブな入力待ちは必要ありません。そんなときは syscall パッケージを利用します。

syscallパッケージ

const ioctlReadTermios = 0x5401  // syscall.TCGETS

// IsTerminal returns true if the given file descriptor is a terminal.
func IsTerminal() bool {
    fd := syscall.Stdin
    var termios syscall.Termios
    _, _, err := syscall.Syscall6(syscall.SYS_IOCTL, uintptr(fd), ioctlReadTermios, uintptr(unsafe.Pointer(&termios)), 0, 0, 0)
    return err == 0
}

上記の IsTerminal 関数が true で返ってくる場合、インタラクティブにデータを受け取る場合になります。処理としては、/dev/stdin のファイルディスクリプタ値である syscall.Stdin が読み込みをしているかを判定しています。

ただ、この実装は Linux に依存しており、他のPlatformで利用するには ioctlReadTermios の値を適宜変更しなければいけません。これを自分で実装するのはちょい面倒なので、既に実装がされているものを利用します。

terminalパッケージ

golang.orgに golang.org/x/crypto/ssh/terminal というパッケージが存在していて、ここにあるIsTerminalという関数が先ほどの各Platformの要件を満たしています。

go get -u golang.org/x/crypto/ssh/terminal

import "golang.org/x/crypto/ssh/terminal"

func main() {
    // ...
    if terminal.IsTerminal(syscall.Stdin) {
        // Do something ...
    }
    // ...
}

全体の実装

package main

import (
    "fmt"
    "io/ioutil"
    "os"
    "syscall"

    "golang.org/x/crypto/ssh/terminal"
)

func main() {
    if terminal.IsTerminal(syscall.Stdin) {
        // Execute: go run main.go
        fmt.Print("Type something then press the enter key: ")
        var stdin string
        fmt.Scan(&stdin)
        fmt.Printf("Result: %s\n", stdin)
        return
    }

    // Execute: echo "foo" | go run main.go
    body, err := ioutil.ReadAll(os.Stdin)
    if err != nil {
        panic(err)
    }
    fmt.Printf("Result: %s\n", string(body))
}