make full screenshot with selenium in python, mobile and web

screenshot.py

import os
import pathlib
from io import BytesIO
from urllib.parse import quote

from PIL import Image


# https://gist.github.com/fabtho/13e4a2e7cfbfde671b8fa81bbe9359fb
class Screenshot:
    @staticmethod
    def full_screenshot_with_scroll(driver, save_path):
        # initiate value

        save_path = save_path.with_suffix("png") if not save_path.match("*.png") else save_path
        img_li = []  # to store image fragment
        offset = 0  # where to start

        # js to get height
        height = driver.execute_script("return Math.max(" "document.documentElement.clientHeight, window.innerHeight);")

        # js to get the maximum scroll height
        # Ref--> https://stackoverflow.com/questions/17688595/finding-the-maximum-scroll-position-of-a-page
        max_window_height = driver.execute_script(
            "return Math.max("
            "document.body.scrollHeight, "
            "document.body.offsetHeight, "
            "document.documentElement.clientHeight, "
            "document.documentElement.scrollHeight, "
            "document.documentElement.offsetHeight);"
        )

        # looping from top to bottom, append to img list
        # Ref--> https://gist.github.com/fabtho/13e4a2e7cfbfde671b8fa81bbe9359fb
        while offset < max_window_height:
            # Scroll to height
            driver.execute_script(f"window.scrollTo(0, {offset});")
            img = Image.open(BytesIO((driver.get_screenshot_as_png())))
            img_li.append(img)
            offset += height

        # In case it is not a perfect fit, the last image contains extra at the top.
        # Crop the screenshot at the top of last image.
        extra_height = offset - max_window_height
        if extra_height > 0 and len(img_li) > 1:
            pixel_ratio = driver.execute_script("return window.devicePixelRatio;")
            extra_height *= pixel_ratio
            last_image = img_li[-1]
            width, height = last_image.size
            box = (0, extra_height, width, height)
            img_li[-1] = last_image.crop(box)

        # Stitch image into one
        # Set up the full screen frame
        img_frame_height = sum([img_frag.size[1] for img_frag in img_li])
        img_frame = Image.new("RGB", (img_li[0].size[0], img_frame_height))
        offset = 0
        for img_frag in img_li:
            img_frame.paste(img_frag, (0, offset))
            offset += img_frag.size[1]
        img_frame.save(save_path)

    @staticmethod
    def save_current_page_screenshot(driver, is_mobile=False, file_name=None, folder=None):
        if file_name is None:
            path = quote(driver.current_url, safe="")
            browser_name = driver.capabilities["browserName"]
            file_name = f"{path}_{browser_name}_screenshot.png"
        else:
            file_name = f"{file_name}.png"

        path = pathlib.Path(__file__)
        file_path = path.parents[1] / "screenshots"
        if folder is not None:
            file_path = file_path / folder / file_name
        else:
            file_path = file_path / file_name

        if is_mobile:
            # PIL will fail if a path or even a relative path is provided,
            # because it expects that the file is already created.
            Screenshot.full_screenshot_with_scroll(driver, file_path)
        else:
            height = driver.execute_script("return document.body.scrollHeight")
            width = driver.get_window_size().get("width", 1024)
            driver.set_window_size(width, height + 100)
            driver.save_screenshot(str(file_path))