リダイレクションやパイプと、端末の入力を同時に行う方法

お題

リダイレクションやパイプからデータを読み込んだ上で端末の入力を読み込む、
そんなプログラムを書こうとしてはまったので整理。

イメージ:

$ echo hoge foo bar | ./to_upper_case_and_xxx.py
HOGE FOO BAR
input:  <- 端末からなにか入力

この場合、次のように書きたくなります。

import sys

# 標準入力をすべて読み込む
for line in sys.stdin:
    sys.stdout.write(line.upper())

# 端末の入力を読み込む
result = raw_input('input: ')
print '>> ' + result

しかし、EOFErrorが発生してうまくいきません。

$ echo hoge | ./raw_input_test.py 
...
EOFError: EOF when reading a line

raw_input()が内部でsys.stdinを参照しているので、読みきってると動かないわけですね。
それに読み込み元がリダイレクションやパイプになっているため、このままではインタラクティブな操作はできません。

ではどうするか

	実行例	標準入力
通常	./program.py	/dev/tty
パイプ	echo hoge ｜ ./program.py	echoの出力
リダイレクション	./program.py < path/to/file	/path/to/file

通常、標準入力は/dev/ttyになっており、端末の入力はこの特殊なファイルから読み込んでいます。
これが、パイプやリダイレクションを利用した場合、別プログラムの出力やファイルに差し替わります。

/dev/ttyをオープンすれば端末からの入力を受け取れるので、sys.stdinに割り当ててみましょう。

修正版

というわけで、次のように変更します。

import os
import sys

# 標準入力をすべて読み込む
for line in sys.stdin:
    sys.stdout.write(line.upper())

# TTY (/dev/tty) をオープンして、標準入力に割り当てる
sys.stdin = file('/dev/tty')

# TTYを割り当てたので、端末の入力を読み込める
result = raw_input('input: ')
print '>> ' + result

今度はちゃんと動作しますね。

もちろん、raw_input()を使わずに直接TTYから読み込んでもOKです。

tty = file(os.ctermid())
sys.stdout.write('input: ')
result = tty.readline()
print '>> ' + result

[python][unix] リダイレクションやパイプの有無を識別する方法

リダイレクションやパイプを使用している場合は読み込み、使用してない場合はヘルプを表示するなど、
切り替えたい場合は次のように書くのがいいのかな。

import os
import sys

if os.isatty(os.sys.stdin.fileno()):
    print_usage()
    sys.exit(1)

# 標準入力をすべて読み込む
for line in sys.stdin:
    sys.stdout.write(line)