python webpage screenshot

python-webpage-screenshot

#!/usr/bin/python
# -*- coding: utf-8 -*-

import uuid
import time
import optparse
import json
import urllib2
import os
import logging

from selenium import webdriver
from pyvirtualdisplay import Display
from PIL import Image


SAVE_PATH = "img"
DRIVERS = ["firefox", "phantom"]
DRIVER = "firefox"
WAIT_INTERVAL = 2
URL = None
BROWSER = None
IMAGE_QUALITY = 60

def init():
	global URL, DRIVER, DRIVERS, WAIT_INTERVAL, SAVE_PATH, IMAGE_QUALITY
	parser = optparse.OptionParser(usage=" ")

	parser.add_option("-u", "--url", dest="url", metavar="http://www.google.com")
	parser.add_option("-d", "--driver", dest="driver", metavar="firefox or phantom")
	parser.add_option("-i", "--interval", dest="interval", metavar="2")
	parser.add_option("-p", "--path", dest="save_path", metavar="img")
	parser.add_option("-q", "--quality", dest="quality", metavar="90")

	(options, args) = parser.parse_args()

	if options.url is not None:
		URL = options.url
	else:
		parser.print_help()
		exit()

	if options.driver is not None:
		if options.driver in DRIVERS:
			DRIVER = options.driver
		else:
			parser.print_help()
			exit()

	if options.interval is not None:
		try:
			WAIT_INTERVAL = float(options.interval)
		except Exception, msg:
			pass

	if options.save_path is not None:
		SAVE_PATH = options.save_path

	if options.quality is not None:
		if options.quality.isdigit():
			quality = int(options.quality)
			if quality < 1:
				quality = 1

			if quality > 100:
				quality = 100

			IMAGE_QUALITY = quality

	if not os.path.exists(SAVE_PATH):
		try:
			os.makedirs(SAVE_PATH)
		except Exception, msg:
			message("error", str(msg))
			exit()


def message(status, message):
	data = {"status": status, "message": message}
	print json.dumps(data)


def create_display():
	try:
		display = Display(visible=0, size=(1280, 1024))
		display.start()
		logging.info("Virtual Display Started")
	except:
		pass


def create_browser():
	global BROWSER

	if DRIVER == "firefox":
		try:
			BROWSER = webdriver.Firefox()
		except Exception, msg:
			message("error", str(msg))
			exit()

	if DRIVER == "phantom":
		try:
			BROWSER = webdriver.PhantomJS()
		except Exception, msg:
			message("error", str(msg))
			exit()


def check_url():
	global URL
	try:
		req = urllib2.urlopen(URL, timeout=20)
		if req.getcode() != 200:
			message("error", "url could not open")
			exit()
	except Exception, msg:
		message("error", str(msg))
		exit()


def take_screenshot():
	global URL, BROWSER, SAVE_PATH, WAIT_INTERVAL, IMAGE_QUALITY

	SAVE_FILE_TEMP = "%s/%s_tmp.png" % (SAVE_PATH, uuid.uuid1())
	SAVE_FILE = "%s/%s.jpg" % (SAVE_PATH, uuid.uuid1())
	
	try:
		BROWSER.get(URL)

		body = BROWSER.find_element_by_css_selector('body')
		body.click()

		time.sleep(WAIT_INTERVAL)

		BROWSER.get_screenshot_as_file(SAVE_FILE_TEMP)

		BROWSER.close()

		im = Image.open(SAVE_FILE_TEMP)

		im.save(SAVE_FILE, "JPEG", quality=IMAGE_QUALITY)

		os.remove(SAVE_FILE_TEMP)

	except Exception, msg:
		message("error", str(msg))
		exit()

	message("ok", SAVE_FILE)

if __name__ == "__main__":
	init()
	check_url()
	create_display()
	create_browser()
	take_screenshot()

requires selenium, pyvirtualdisplay, PIL

tested on python 2.7.5

usage: python app.py --url "http://en.wikipedia.org/wiki/Python_(programming_language)" --driver firefox --interval 2 --path imgs --quality 70