PDF Add Page Number

main.py

# !/usr/bin/env python3
# coding: utf-8
"""
*PDF Add Page Number*

[Principle]
a. create a layer pdf with page numbers
b. merge the origin pdf with the layer pdf
"""
import sys
from pathlib import Path

from PyPDF2.pdf import PdfFileReader, PdfFileWriter
from reportlab.lib.units import mm
from reportlab.pdfgen import canvas


class PDF(object):

    def __init__(self, file: str):
        self.origin = PdfFileReader(Path(file).open('rb'), strict=False)

        self.layer_path = Path(file).with_stem('temp')
        self.output_path = Path(file).with_stem(f'{Path(file).stem}_with_page_number')

    @property
    def number_pages(self):
        return self.origin.getNumPages()

    @property
    def width(self):
        return self.origin.getPage(0).mediaBox.getWidth()

    @property
    def height(self):
        return self.origin.getPage(0).mediaBox.getHeight()

    @property
    def pagesize(self):
        return self.width, self.height

    def create_page_number_layer(self):
        c = canvas.Canvas(self.layer_path.as_posix(), pagesize=self.pagesize)
        for i in range(self.number_pages):
            c.drawString(float(self.width) // 2, 4 * mm, str(i + 1))
            c.showPage()
        c.save()

    def add_page_number(self):
        output = PdfFileWriter()

        self.create_page_number_layer()
        layer_pdf = PdfFileReader(self.layer_path.open('rb'), strict=False)

        for i in range(self.number_pages):
            page = self.origin.getPage(i)
            page.mergePage(layer_pdf.getPage(i))
            output.addPage(page)

        output.write(self.output_path.open('wb'))


if __name__ == '__main__':
    if len(sys.argv) > 1 and Path(sys.argv[1]).exists():
        PDF(sys.argv[1]).add_page_number()