Skip to content

A few improvements #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 7 commits into from
Closed

A few improvements #62

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c883788
Replace "if self.logger: ..." with default logger
dotdoom Mar 31, 2016
81460b1
Add time.sleep() in busy-wait loops to relieve CPU
dotdoom Mar 31, 2016
791a576
Propagade cookies to all requests (e.g. images)
dotdoom Mar 31, 2016
9fbc360
Do not abort on rendering timeout
dotdoom Mar 31, 2016
96f0a9c
Allow overriding User-Agent header
dotdoom Mar 31, 2016
dc22a12
Resize the window only after "delay" is reached
dotdoom Mar 31, 2016
98c8916
Allow rendering only specific element
dotdoom Mar 31, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 65 additions & 28 deletions webkit2png/webkit2png.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
# - Add QTcpSocket support to create a "screenshot daemon" that
# can handle multiple requests at the same time.

import logging
import time
import os

Expand All @@ -31,6 +32,8 @@
from PyQt4.QtWebKit import *
from PyQt4.QtNetwork import *

defaultLogger = logging.getLogger("webkit2png")

# Class for Website-Rendering. Uses QWebPage, which
# requires a running QtGui to work.
class WebkitRenderer(QObject):
Expand Down Expand Up @@ -60,19 +63,21 @@ def __init__(self,**kwargs):
self.scaleToHeight = kwargs.get('scaleToHeight', 0)
self.scaleRatio = kwargs.get('scaleRatio', 'keep')
self.format = kwargs.get('format', 'png')
self.logger = kwargs.get('logger', None)
self.logger = kwargs.get('logger', defaultLogger)

# Set this to true if you want to capture flash.
# Not that your desktop must be large enough for
# fitting the whole window.
self.grabWholeWindow = kwargs.get('grabWholeWindow', False)
self.elementSelector = kwargs.get('elementSelector', None)
self.renderTransparentBackground = kwargs.get('renderTransparentBackground', False)
self.ignoreAlert = kwargs.get('ignoreAlert', True)
self.ignoreConfirm = kwargs.get('ignoreConfirm', True)
self.ignorePrompt = kwargs.get('ignorePrompt', True)
self.interruptJavaScript = kwargs.get('interruptJavaScript', True)
self.encodedUrl = kwargs.get('encodedUrl', False)
self.cookies = kwargs.get('cookies', [])
self.userAgent = kwargs.get('userAgent', 'WebKit2HTML/1.0')

# Set some default options for QWebPage
self.qWebSettings = {
Expand Down Expand Up @@ -149,6 +154,8 @@ def __init__(self, parent):
"""
QObject.__init__(self)

self.__loading_result = None

# Copy properties from parent
for key,value in parent.__dict__.items():
setattr(self,key,value)
Expand All @@ -173,7 +180,7 @@ def __init__(self, parent):
# Create and connect required PyQt4 objects
self._page = CustomWebPage(logger=self.logger, ignore_alert=self.ignoreAlert,
ignore_confirm=self.ignoreConfirm, ignore_prompt=self.ignorePrompt,
interrupt_js=self.interruptJavaScript)
interrupt_js=self.interruptJavaScript, userAgent=self.userAgent)
self._page.networkAccessManager().setProxy(proxy)
self._view = QWebView()
self._view.setPage(self._page)
Expand Down Expand Up @@ -216,16 +223,26 @@ def render(self, res):
on the value of 'grabWholeWindow' is drawn into a QPixmap
and postprocessed (_post_process_image).
"""
self._load_page(res, self.width, self.height, self.timeout)
self._load_page(res, self.timeout)
# Wait for end of timer. In this time, process
# other outstanding Qt events.
if self.wait > 0:
if self.logger: self.logger.debug("Waiting %d seconds " % self.wait)
self.logger.debug("Waiting %d seconds " % self.wait)
waitToTime = time.time() + self.wait
while time.time() < waitToTime:
time.sleep(0.1)
if QApplication.hasPendingEvents():
QApplication.processEvents()

# Set initial viewport (the size of the "window")
size = self._page.mainFrame().contentsSize()
self.logger.debug("contentsSize: %s x %s", size.width(), size.height())
if self.width > 0:
size.setWidth(self.width)
if self.height > 0:
size.setHeight(self.height)
self._window.resize(size)

if self.renderTransparentBackground:
# Another possible drawing solution
image = QImage(self._page.viewportSize(), QImage.Format_ARGB32)
Expand All @@ -241,6 +258,26 @@ def render(self, res):
painter.setBackgroundMode(Qt.TransparentMode)
self._page.mainFrame().render(painter)
painter.end()
elif self.elementSelector:
elementAndSize = self.elementSelector(
self._page.mainFrame().documentElement())
if isinstance(elementAndSize, tuple):
(element, size) = elementAndSize
else:
(element, size) = (elementAndSize,
elementAndSize.geometry().size())
if size.isEmpty():
raise RuntimeError("Selected element is empty")
else:
self.logger.debug("Selected element size: %d x %d",
size.width(),
size.height())
self._window.resize(size)
image = QImage(size, QImage.Format_ARGB32)
image.fill(QColor(255,0,0,0).rgba())
painter = QPainter(image)
element.render(painter)
painter.end()
else:
if self.grabWholeWindow:
# Note that this does not fully ensure that the
Expand All @@ -253,7 +290,7 @@ def render(self, res):

return self._post_process_image(image)

def _load_page(self, res, width, height, timeout):
def _load_page(self, res, timeout):
"""
This method implements the logic for retrieving and displaying
the requested page.
Expand Down Expand Up @@ -281,7 +318,11 @@ def _load_page(self, res, width, height, timeout):
qtUrl = QUrl(url)

# Set the required cookies, if any
self.cookieJar = CookieJar(self.cookies, qtUrl)
urlWithoutPath = QUrl(qtUrl)
# Reset the path to root: QT will only serve the cookie
# to the child components of the initial path.
urlWithoutPath.setPath("/")
self.cookieJar = CookieJar(self.cookies, urlWithoutPath)
self._page.networkAccessManager().setCookieJar(self.cookieJar)

# Load the page
Expand All @@ -291,25 +332,17 @@ def _load_page(self, res, width, height, timeout):
self._page.mainFrame().load(qtUrl)

while self.__loading:
time.sleep(0.1)
if timeout > 0 and time.time() >= cancelAt:
raise RuntimeError("Request timed out on %s" % res)
self.logger.warning("Request timed out on %s" % res)
break
while QApplication.hasPendingEvents() and self.__loading:
QCoreApplication.processEvents()

if self.logger: self.logger.debug("Processing result")
self.logger.debug("Processing result")

if self.__loading_result == False:
if self.logger: self.logger.warning("Failed to load %s" % res)

# Set initial viewport (the size of the "window")
size = self._page.mainFrame().contentsSize()
if self.logger: self.logger.debug("contentsSize: %s", size)
if width > 0:
size.setWidth(width)
if height > 0:
size.setHeight(height)

self._window.resize(size)
self.logger.warning("Failed to load %s" % res)

def _post_process_image(self, qImage):
"""
Expand Down Expand Up @@ -341,15 +374,15 @@ def _on_load_started(self):
"""
Slot that sets the '__loading' property to true
"""
if self.logger: self.logger.debug("loading started")
self.logger.debug("loading started")
self.__loading = True

# Eventhandler for "loadFinished(bool)" signal
def _on_load_finished(self, result):
"""Slot that sets the '__loading' property to false and stores
the result code in '__loading_result'.
"""
if self.logger: self.logger.debug("loading finished with result %s", result)
self.logger.debug("loading finished with result %s", result)
self.__loading = False
self.__loading_result = result

Expand All @@ -359,7 +392,7 @@ def _on_ssl_errors(self, reply, errors):
Slot that writes SSL warnings into the log but ignores them.
"""
for e in errors:
if self.logger: self.logger.warn("SSL: " + e.errorString())
self.logger.warning("SSL: " + e.errorString())
reply.ignoreSslErrors()


Expand All @@ -369,19 +402,23 @@ def __init__(self, **kwargs):
Class Initializer
"""
super(CustomWebPage, self).__init__()
self.logger = kwargs.get('logger', None)
self.logger = kwargs.get('logger', defaultLogger)
self.ignore_alert = kwargs.get('ignore_alert', True)
self.ignore_confirm = kwargs.get('ignore_confirm', True)
self.ignore_prompt = kwargs.get('ignore_prompt', True)
self.interrupt_js = kwargs.get('interrupt_js', True)
self.userAgent = kwargs['userAgent']

def userAgentForUrl(self, url):
return self.userAgent

def javaScriptAlert(self, frame, message):
if self.logger: self.logger.debug('Alert: %s', message)
self.logger.debug('Alert: %s', message)
if not self.ignore_alert:
return super(CustomWebPage, self).javaScriptAlert(frame, message)

def javaScriptConfirm(self, frame, message):
if self.logger: self.logger.debug('Confirm: %s', message)
self.logger.debug('Confirm: %s', message)
if not self.ignore_confirm:
return super(CustomWebPage, self).javaScriptConfirm(frame, message)
else:
Expand All @@ -398,7 +435,7 @@ def javaScriptPrompt(self, frame, message, default, result):
If the prompt was not cancelled by the user, the implementation should return true and
the result string must not be null.
"""
if self.logger: self.logger.debug('Prompt: %s (%s)' % (message, default))
self.logger.debug('Prompt: %s (%s)' % (message, default))
if not self.ignore_prompt:
return super(CustomWebPage, self).javaScriptPrompt(frame, message, default, result)
else:
Expand All @@ -409,5 +446,5 @@ def shouldInterruptJavaScript(self):
This function is called when a JavaScript program is running for a long period of time.
If the user wanted to stop the JavaScript the implementation should return true; otherwise false.
"""
if self.logger: self.logger.debug("WebKit ask to interrupt JavaScript")
self.logger.debug("WebKit ask to interrupt JavaScript")
return self.interrupt_js