仮想ディスプレイとSeleniumでスクショを撮影する方法

しんぺい

こんにちは!しんぺいです。

Pythonのselenium便利ですよねー。
ソースコードパクったり、テストを自動化したり、情報解析したりと、出来ることの幅がグッと広がります。

今回はそんなSeleniumを仮想ディスプレイで動かして、スクリーンショットを撮影する方法を紹介します!

pyvirtualdisplayとSeleniumのインストール

まずは必要なものをインストールします。
Python側で必要なのは、仮想ディスプレイを起動させるpyvirtualdisplayと、Seleniumですね。

$ pip install pyvirtualdisplay selenium

また、mac上でx11を表示するためにXQuartzもインストールして、起動させておきます。
XQuartzがインストールされていないと、こんなエラーが発生する場合があります。

pyvirtualdisplay.abstractdisplay.XStartTimeoutError: Failed to start X on display ":1001" (xdpyinfo check failed).

Pythonコードを実装する

あとはサクッとPythonコードを実装しちゃいましょう。

#!/usr/bin/env python

from pyvirtualdisplay import Display
from selenium import webdriver

import time

UA = 'Mozilla/5.0 (iPhone; CPU iPhone OS 10_1_2 like Mac OS X) AppleWebKit/354.4.6 (KHTML, like Gecko) Mobile/14D19 Safari/354.4.6 chrome/3.3.2'

display = Display(visible=0, size=(800, 600))
display.start()

profile = webdriver.FirefoxProfile()
profile.set_preference("general.useragent.override", UA)
driver = webdriver.Firefox(firefox_profile=profile, executable_path=r'/usr/local/bin/geckodriver')
driver.set_window_size(900, 1600)
driver.implicitly_wait(30)
driver.set_page_load_timeout(30)

driver.get('http://www.google.com')
time.sleep(5)
driver.save_screenshot('{0}.png'.format("google"))
driver.quit()

display.stop()

順番に見ていきます。
最初のUAには、あとでプロファイルのユーザーエージェントを上書きするためのものを入れています。
ユーザーエージェントが変更できる、ということです。

UA = 'Mozilla/5.0 (iPhone; CPU iPhone OS 10_1_2 like Mac OS X) AppleWebKit/354.4.6 (KHTML, like Gecko) Mobile/14D19 Safari/354.4.6 chrome/3.3.2'

次の2行で仮想ディスプレイを立ち上げています。
特に難しいことはしていません。

display = Display(visible=0, size=(800, 600))
display.start()

次の2行でFirefoxのプロファイルを取得し、ユーザーエージェントを書き換えています。

profile = webdriver.FirefoxProfile()
profile.set_preference("general.useragent.override", UA)

次の行でユーザーエージェントを書き換えたプロファイルをFirefoxに設定し、ドライバーを起動しています。
ドライバーが置かれている場所は適宜変更してください。

driver = webdriver.Firefox(firefox_profile=profile, executable_path=r'/usr/local/bin/geckodriver')

implicitly_waitでドライバーの要素が見つかるまでの待機時間を設定します。
見つからなかった場合にはExceptionが発生します。

set_page_load_timeoutでは、ページが完全にロードされるまでの待ち時間を指定しています。
待機時間を超えてもロードが完了しなかった場合にはExceptionが発生します。

driver.set_window_size(900, 1600)
driver.implicitly_wait(30)
driver.set_page_load_timeout(30)

あとはスクリーンショットを撮りたいURLを指定して、撮影します。
画像の置き場所は特に指定していないので同階層に保存されますが、運用面などを考えるのであれば、path.dirname(path.realpath(__file__))などを挟んだ方がいいかもしれませんね。

driver.get('http://www.google.com')
time.sleep(5)
driver.save_screenshot('{0}.png'.format("google"))

最後に仮想ディスプレイ、ブラウザを閉じる処理を忘れないようにしましょう!

driver.quit()

display.stop()

以上です!