browser screenshot with selenium

selenium_driver.py

from PIL import Image
from selenium import webdriver
from selenium.webdriver.chrome.options import Options


class Driver:
    """ selenium web driver instance """

    def __init__(self):
        self.chrome = None
        self.initial()

    def initial(self):
        if self.chrome is None:
            options = Options()
            options.headless = True
            self.chrome = webdriver.Chrome(options=options)

    def save_screenshot(self, url, filename, div_id=None):
        """
        :param url: 访问地址
        :param filename: 文件图片保存地址
        :param div_id: dom ID
        :return:
        """
        self.chrome.get(url=url)
        self.chrome.maximize_window()
        self.chrome.save_screenshot(filename)
        if div_id is not None:
            elem = self.chrome.find_element_by_id(div_id)
            location_x, location_y = elem.location['x'], elem.location['y']
            size_w, size_h = elem.size['width'], elem.size['height']
            box = (location_x, location_y, location_x + size_w, location_y + size_h)
            self.img_crop(filename, box)

    def close(self):
        """ 关闭模拟浏览器 """
        if self.chrome is not None:
            self.chrome.close()

    @staticmethod
    def img_crop(filename, box, output=None):
        """
        对图片进行裁剪
        :param filename: 原始图片地址
        :param box: 裁剪区域
        :param output: 处理后图片保存地址
        :return:
        """
        if output is None:
            output = filename
        with Image.open(filename) as img:
            img = img.crop(box)
            img.save(output)


chrome = Driver()