WeasyPrint/weasyprint/pdf.py

# coding: utf8
"""
    weasyprint.pdf
    --------------

    :copyright: Copyright 2011-2012 Simon Sapin and contributors, see AUTHORS.
    :license: BSD, see LICENSE for details.

"""

from __future__ import division, unicode_literals


class PDF(object):
    def __init__(self, bytesio, links, destinations):
        bytesio.seek(0)
        self.lines = bytesio.readlines()
        self.outlines = []
        self.xref = []
        self.trailer = []
        self.size = None
        self.info = None
        self.objects = {}
        self.active = None
        self.numbers = []
        self.added_numbers = []
        self.pages = []

        for line in self.lines:
            if line.endswith(b' obj\n'):
                self.active = 'object'
                number = int(line.split()[0])
                self.numbers.append(number)
            elif line == b'xref\n':
                self.active = 'xref'
            elif line == b'trailer\n':
                self.active = 'trailer'

            if line.endswith(b'/Type /Page\n'):
                self.pages.append(number)

            if self.active == 'size':
                self.size = int(line)
                self.active = None
            elif self.active == 'xref':
                self.xref.append(line)
            elif self.active == 'trailer':
                self.trailer.append(line)
            elif self.active == 'object':
                if number not in self.objects:
                    self.objects[number] = []
                self.objects[number].append(line)

            if line == b'startxref\n':
                self.active = 'size'
            elif line == b'endobj\n':
                self.active = None

        for i, line in enumerate(self.trailer):
            if '/Info' in line:
                self.info = int(line.rsplit()[-3])

        for i, line in enumerate(self.objects[self.info]):
            if '/Creator' in line:
                pre = line.split('/Creator')[0]
                new_line = '%s/Creator (%s)\n' % (pre, 'WeasyPrint')
                self.objects[self.info][i] = new_line
                offset_size = len(new_line) - len(line)
                self.replace_xref_size(self.info, offset_size)

        for pdf_page_number, link_page in zip(self.pages, links):
            annot_numbers = []
            for link, x1, y1, x2, y2 in link_page:
                if link.startswith('#'):
                    destination = destinations.get(link[2:])
                    if not destination:
                        continue
                    annot_numbers.append(self.add_object(b"""<<
/Type /Annot
/Subtype /Link
/Rect [%f %f %f %f]
/A <<
/Type /Action
/S /GoTo
/D [%d /XYZ %d %d 1]
>>
>>""" % (x1, y1, x2, y2, destination[0], destination[1], destination[2])))
                else:
                    annot_numbers.append(self.add_object(b"""<<
/Type /Annot
/Subtype /Link
/Rect [%f %f %f %f]
/A <<
/Type /Action
/S /URI
/URI (%s)
>>
>>""" % (x1, y1, x2, y2, link)))

            string = '/Annots [%s]\n' % ' '.join(
                '%d 0 R' % number for number in annot_numbers)
            self.objects[pdf_page_number].insert(12, string)
            self.replace_xref_size(pdf_page_number, len(string))

        for i, line in enumerate(self.trailer):
            if '/Size' in line:
                pre = line.split('/Size')[0]
                new_line = '%s/Size %s\n' % (
                    pre, '%s' % (len(self.added_numbers) + len(self.numbers)))
                self.trailer[i] = new_line

        for line in self.lines:
            if line.endswith(b' obj\n'):
                self.active = 'object'
                number = int(line.split()[0])
                self.outlines.extend(self.objects[number])
            elif line == b'xref\n':
                self.active = 'xref'
                for added_number in self.added_numbers:
                    self.outlines.extend(self.objects[added_number])
                self.outlines.extend(self.xref)
            elif line == b'trailer\n':
                self.active = 'trailer'

            if self.active == 'size':
                self.outlines.append('%d\n' % self.size)
                self.active = None
            elif self.active == 'xref':
                pass
            elif self.active == 'object':
                pass
            elif self.active == 'trailer':
                if self.trailer:
                    self.outlines.extend(self.trailer)
                    self.trailer = None
            else:
                self.outlines.append(line)

            if line == b'endobj\n':
                self.active = None
            elif line == b'startxref\n':
                self.active = 'size'
            elif line == b'endobj\n':
                self.active = None

    def add_object(self, text):
        next_number = len(self.numbers) + 1
        text = '%d 0 obj\n%s\nendobj' % (next_number, text)
        last_size = int(self.xref[self.numbers[-1] + 2].split()[0].lstrip('0'))
        last_object_size = len(''.join(self.objects[self.numbers[-1]]))
        self.xref.append('%010d 00000 n \n' % (last_size + last_object_size))
        self.numbers.append(next_number)
        self.added_numbers.append(next_number)
        self.objects[next_number] = [line + '\n' for line in text.split('\n')]
        self.size += len(text) + 1
        self.xref[1] = '0 %d\n' % (next_number + 1)
        return next_number

    def replace_xref_size(self, number, offset_size):
        index = self.numbers.index(number)
        for next_number in self.numbers[index + 1:len(self.numbers)]:
            out = self.xref[next_number + 2]
            old_size, content = out.split(' ', 1)
            old_size = int(old_size.lstrip('0'))
            old_size += offset_size
            self.xref[next_number + 2] = '%010d %s' % (old_size, content)
        self.size += offset_size

    def write(self, target):
        for outline in self.outlines:
            target.write(outline)