python+selenium+webdriver自动化截取整站截图

背景介绍

最近在为部门编写一个网站截图工具，将输入网址自动截取全站截图。

方案思考

python+selenium+webdriver

github：

https://github.com/iamjasonchoi/Screenshot

主要代码如下:
————————————————

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
import requests
from bs4 import BeautifulSoup
import time
from tld import get_fld
import os

i = 0
# 目前暂时先用一个计数来表示图片的文件名

def getScreenShot(url):
    '''
    此函数用于根据url访问网页，并截图到本地文件夹
    '''
    global i  # 全局变量
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-gpu')  # 设置浏览器参数

    chrome_options.add_argument("--window-size=1440,1024")
    chrome_options.add_argument("--hide-scrollbars")

    brower = webdriver.Chrome(chrome_options=chrome_options)  # 创建一个chrome的webdrive
    brower.maximize_window()  # 设置全屏截图
    js_height = "return document.body.clientHeight"

    path = "/Users/xxx/PycharmProjects/CutwebV2/venv/pic/" + get_fld(url) + "/"
    #path = "C:/pic/" + urlparse.urlparse(url).netloc + "/"
    # 判断路径是否存在
    isExists = os.path.exists(path)

    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数

        os.makedirs(path)

    else:
        path = "/Users/xxx/PycharmProjects/CutwebV2/venv/pic/" + get_fld(url) + "/"
        #path = "C:/pic/" + urlparse.urlparse(url).netloc + "/"

    picName = path + get_fld(url) + "_" + str(i) + ".png"  # 指定保存文件的文件名

    brower.get(url)  # 获取url
    i += 1

    try:
        brower.get(url)
        k = 1
        height = brower.execute_script(js_height)
        # 获取整张网页的截图，而不只是当前屏幕范围内的

        while True:
            if k * 500 < height:
                js_move = "window.scrollTo(0,{})".format(k * 500)
                print("============working============")
                brower.execute_script(js_move)
                time.sleep(0.2)
                height = brower.execute_script(js_height)
                k += 1
            else:
                break
        scroll_width = brower.execute_script('return document.body.parentNode.scrollWidth')
        scroll_height = brower.execute_script('return document.body.parentNode.scrollHeight')
        brower.set_window_size(scroll_width, scroll_height)
        brower.get_screenshot_as_file(picName)
        time.sleep(0.1)
    except Exception as e:
        print(picName, e)

    brower.save_screenshot(picName)  # 保存截图
    brower.close()  # 关闭webdriver


def getHTMLText(url):
    '''
    此函数用于获取网页的html文档
    '''
    try:
        # 获取服务器的响应内容，并设置最大请求时间为6秒
        res = requests.get(url, timeout=6)
        # 判断返回状态码是否为200
        res.raise_for_status()
        # 设置该html文档可能的编码
        res.encoding = res.apparent_encoding
        # 返回网页HTML代码
        return res.text
    except:
        return '产生异常'


def main():
    '''
    主函数
    '''
    # 目标网页，这个可以换成一个你想要的网站
    #url = raw_input("请输入网址【形如http://www.xxx.com】: ");
    #url = "' " + url + "'"
    #test
    #url = 'http://tjhglngy.com/'

    f = open("urls.txt", "r")
    urlList = f.readlines()
    for url in urlList:
        try:
            demo = getHTMLText(url)  # 获取html代码

            # 解析HTML代码
            soup = BeautifulSoup(demo, 'html.parser')

            # 模糊搜索HTML代码的所有包含href属性的<a>标签
            a_labels = soup.find_all('a', attrs={'href': True})

            hash_set = set()  # 创建一个哈希表，避免对同一个网站重复截图，每当访问一个网站后，将网址加入到哈希表中，如果哈希表中已经有这个网址，则不再截图
            # 获取所有<a>标签中的href对应的值，即超链接
            for a in a_labels:
                s = url + a.get('href')
                if s not in hash_set:  # 如果这个超链接不在哈希表中
                    hash_set.add(s)  # 添加这超链接到哈希表
                    getScreenShot(s)  # 获取这个网站的截图
            return url + '站点截图完成！'
        except:
            return '产生异常了'

if __name__ == "__main__":
     main()