webpage2img/webpage2img.py

#!/usr/bin/python3
#
# Copyright 2017 Weber Yann
#
# This file is part of webpage2img.
#
# webpage2img is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# any later version.
#
# webpage2img is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with webpage2img.  If not, see <http://www.gnu.org/licenses/>.
#

import sys
import argparse
from PySide import QtWebKit, QtCore, QtGui

#self._app = QtGui.QApplication(sys.argv)

class WebPage2Img(object):
    
    def __init__(self, uri):
        self._out_fname = None
        self.uri = uri
        self._app = QtGui.QApplication(sys.argv)
        #After this point we are fucked up because the app catch
        #all exceptions :/
        self.wxview = QtWebKit.QWebView()


    def png(self, filename):
        self.__output(filename, self._png)

    def pdf(self, filename):
       self.__output(filename, self._pdf)
    
    def __output(self, filename, slot):
        self._out_fname = filename
        self.wxview.loadFinished.connect(slot)
        print('PNG rendering to : %s' % self._out_fname)
        self.__load()

    ##@brief PNG rendering function
    #@param ok : no idea, sent by connect
    #@utils.ExitingExptSlot(self._app, "bool")
    def _png(self, ok):
        print('_png')
        self.wxview.page().setViewportSize(self.wxview.page().currentFrame().contentsSize())
        image = QtGui.QImage(self.wxview.page().viewportSize(),QtGui.QImage.Format_ARGB32)
        painter = QtGui.QPainter(image)
        self.wxview.page().mainFrame().render(painter)
        painter.end()
        image.save(self._out_fname)
        self.exit()

    def _pdf(self, ok):
        print('_pdf')
        printer = QtGui.QPrinter()
        printer.setOutputFileName(self._out_fname)
        printer.setOutputFormat(QtGui.QPrinter.PdfFormat)
        printer.setPageSize(QtGui.QPrinter.A4)
        printer.setOrientation(QtGui.QPrinter.Landscape)        
        #printer.setFullPage(True)
        self.wxview.page().mainFrame().print_(printer)
        self.exit()

        
    def exit(self, retcode = 0):
        print('Exiting.')
        self._app.exit(retcode)
        exit(retcode)
        

    def __load(self):
        if self._out_fname is None:
            raise Exception('Do not run if no out set')
        self.wxview.load(QtCore.QUrl(self.uri))
        sys.exit(self._app.exec_())


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Convert a webpage to an \
image.')
    parser.add_argument('-f', '--format',  dest='fmt', metavar='pdf|png', type=str,
        default='pdf', help='Choose output format')
    parser.add_argument('-o', '--output', metavar='FILENAME', type=str,
        default='out.pdf', help='Output file')
    parser.add_argument('uri', metavar='URI', type=str)
    args = parser.parse_args()

    wp = WebPage2Img(args.uri)
    if args.fmt.lower() == 'png':
        wp.png(args.output)
    elif args.fmt.lower() == 'pdf':
        wp.pdf(args.output)
    else:
        print('Unknown format "%s"\n' % args.fmt)
        parser.print_help()
        wp.exit(1)