cgiからmatplotlibで画像を表示させる

簡単な部内サイトなどを作る時、matplotlibで書いた画像を動的に表示させたい場合cgi スクリプトで動かしたい場合があると思う。 matplotlibの標準出力のサンプルを参考に、レスポンスヘッダを付加してすれば良いと思って、

import sys
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

plt.plot([1, 2, 3])

print('Content-type: image/png')
print()

plt.savefig(sys.stdout.buffer, format='png')

のように書いたところ、エラーで表示されない。Apacheのログをみると

malformed header from script 'image.cgi': Bad header: \x89PNG

となっていた。ヘッダが悪い？と思っていろいろ調べてみると、stackoverflowにそれっぽい回答が。print関数の出力が標準出力にうまく返されていない？のが原因らしい。

回答の通り以下のように書き換えたらうまく表示された。

import sys
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

plt.plot([1, 2, 3])

print('Content-type: image/png')
print()
sys.stdout.flush()

plt.savefig(sys.stdout.buffer, format='png'))

もしくは、

import sys
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

plt.plot([1, 2, 3])

sys.stdout.buffer.write(b'Content-type: image/png\n\n') 
plt.savefig(sys.stdout.buffer, format='png'))