Take screenshots of websites and place into PDF

screenshot2pdf.py

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from PIL import Image
import io
import time
import os
import sys
import numpy as np
from fpdf import FPDF

WEBSITES = """https://google.com
https://yahoo.com"""
WEBSITES = WEBSITES.splitlines()
# screenshot size
WIDTH = 800
HEIGHT = 600
# screenshot margin
MARGIN = 10
# screenshot padding
PADDING = 10
options = Options()
options.add_argument('--headless')
options.add_argument('--disable-gpu')
options.add_argument('--no-sandbox')
driver = webdriver.Chrome(options=options, executable_path=os.getcwd() + '/chromedriver')
driver.set_window_size(WIDTH, HEIGHT)
for website in WEBSITES:
    try:
        driver.get(website)
        time.sleep(1)
        # zoom out the image before taking screenshot
        driver.execute_script("document.body.style.zoom='50%'")
        # get the screenshot
        time.sleep(0.2)
        img = Image.open(io.BytesIO(driver.get_screenshot_as_png()))
        img = img.resize((WIDTH, HEIGHT))
    except:
        print('error', website)
    img.save('/tmp/screenshots/' + website.replace('https://', '').replace('http://', '').replace('/', '_') + '.png')
driver.quit()

pdf = FPDF()
images = os.listdir('/tmp/screenshots')
images.sort()
for image in images:
    if not image.endswith('.png'):
        continue
    pdf.add_page()
    pdf.set_font('Arial', 'B', 16)
    pdf.cell(40, 10, image.replace('.png', ''))
    pdf.image('/tmp/screenshots/' + image, x=0, y=20, w=210, h=0, type='png')
pdf.output('/tmp/screenshots.pdf', 'F')