多线程PDF下载器

该PDF下载器支持从本地txt文件中读取和数据库读取两种方式

数据库ip我乱写的，这里集成了三种情况的数据库（本地数据库，需要验证的数据库，不需要验证的数据库）

# -*- coding: utf-8 -*-
"""
此程序用来多线程下载PDF，提供本地txt文件读取下载和数据库读取下载两种方式
"""
import hashlib
import os
import sys
from multiprocessing import Queue
from threading import Thread, Lock

import pymongo
import requests

g_num = 0  # 创建全局变量
lock = Lock()  # 创建全局互斥锁


class DownloadPDF(object):

    def __init__(self):
        # 创建消息队列
        self.q = Queue()

    def run(self):
        self.mkdir()
        self.get_pdf_msg()
        pdf_size = self.q.qsize()  # 总PDF数
        # 开启多线程
        all_th = []
        for i in range(10):
            th = Thread(target=self.download, args=(pdf_size,))
            th.start()
            all_th.append(th)

        for th in all_th:
            th.join()
        print('\n' + '下载完成！')

    def download(self, pdf_size):
        while True:
            try:
                # 获取PDF信息
                each_pdf = self.q.get(timeout=2)  # 设置超时时间2s。这里如果不设置超时，当队列为空时，q.get()会进入阻塞状态
            except:
                # 这里队列已经为空，取不到数据就跳出循环
                break
                # 计数
            global g_num
            lock.acquire()  # 上锁
            g_num += 1
            # 打印下载进度
            sys.stdout.write('\r' + '正在下载：%s / %s' % (g_num, pdf_size))
            path = ''
            pdf_url = ''
            md5_later = ''
            pdf_name = ''  # 如果没有PDF名，置为空
            if is_txt:
                # 选择了txt文件方式，只有一个链接
                pdf_url = each_pdf.strip()
                href_name = pdf_url.split('/')[-1]  # 取链接 / 后的值作为PDF名称
                if '.pdf' in href_name:
                    href_name = href_name.split('.')[0]
                md5_later = self.pdf_name_md5(href_name)  # pdf转码之后的名称
                path = 'PDF/' + md5_later + '.pdf'

            elif not is_txt:
                # 从数据库中读取PDF信息，数据库中包含两个字段(pdf_name, pdf_url)
                pdf_name = each_pdf['pdf_name']
                pdf_url = each_pdf['pdf_url']
                md5_later = self.pdf_name_md5(pdf_name)
                path = 'PDF/' + md5_later + '.pdf'
            # 下载之前先去重
            try:
                # 尝试获取该路径下pdf的size，如果还未下载则置为0
                size = os.path.getsize(path)
            except:
                size = 0
            # 判断已经下载并且size > 0则为成功下载
            if os.path.exists(path) and size != 0:
                # 路径存在说明下载过
                lock.release()  # 释放锁
                continue
            # 下载PDF
            self.download_pdf(pdf_name, each_pdf, path, pdf_url, md5_later)
            lock.release()  # 释放锁

    @staticmethod
    def download_pdf(pdf_name, each_pdf, path, pdf_url, md5_later):
        """
        下载PDF的主要实现程序
        :param pdf_name:
        :param each_pdf:
        :param path:
        :param pdf_url:
        :param md5_later:
        :return:
        """
        try:
            # 开始下载
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)'
                              ' Chrome/76.0.3809.100 Safari/537.36'}
            r = requests.get(url=pdf_url, headers=headers, timeout=(120, 600))  # 设置connect超时2分钟，read超时10分钟
            f = open(path, 'wb')
            f.write(r.content)
            f.close()
            r.close()
            s = requests.session()
            s.keep_alive = False
            if is_insert == 'y':
                # 下载成功，插入数据库
                if not pdf_name:
                    msg = {'url': pdf_url, 'md5_name': md5_later, 'relative_path': path}
                else:
                    msg = {'url': pdf_url, 'origin_name': pdf_name, 'md5_name': md5_later, 'relative_path': path}
                new_collection.insert(msg)
        except:
            # 下载失败
            print('\r' + '下载失败：%s' % str(each_pdf))
            with open('error.txt', 'r') as fr:
                content = fr.readlines()
            if str(each_pdf) not in content:
                with open('error.txt', 'a') as fe:
                    fe.writelines(str(each_pdf) + '\n')

    @staticmethod
    def mkdir():
        """
        创建下载失败文件、去重文件、PDF文件夹
        :return:
        """
        if not os.path.exists('error.txt'):
            with open('error.txt', 'w') as fr:
                fr.write('')
        if not os.path.exists('PDF'):
            os.mkdir('PDF')

    def client_dbs(self):
        """
        连接数据库，获取所有pdf信息
        :return:
        """
        ip = self.choose_ip()

        client = pymongo.MongoClient('mongodb://%s:27017' % ip)
        # 这里因为该数据库开了验证，所以这里进行判断验证（ip我乱写的）
        if ip == '192.168.0.123':
            client['admin'].authenticate('用户名', '密码')  # 数据库验证（用户名密码根据自己数据库修改）
        all_db_names = client.list_database_names()
        # 规范打印数据库名
        self.pr_datas(all_db_names)

        db_name_num = input('\r' + '请选择数据库对应序号：')
        db_name = self.choose_db_or_col(all_db_names, db_name_num, '数据库')

        while True:
            db = client[db_name]
            all_col = db.collection_names()
            if not all_col:
                db_name = input('没有这个数据库，请重新输入：')
            else:
                break
        # 规范打印集合名
        self.pr_datas(all_col)
        col_num = input('\r' + '请选择集合对应的序号：')
        col_name = self.choose_db_or_col(all_col, col_num, '集合')

        collection = db[col_name]
        return collection, db

    @staticmethod
    def choose_ip():
        """
        选择ip
        :return:
        """
        ip = input('请选择数据库：' + '\n' + '1. 本地数据库   2. 123数据库   3. 200数据库' + '\n')
        while True:
            if ip == '1':
                ip = '127.0.0.1'
                break
            elif ip == '2':
                ip = '192.168.0.123'
                break
            elif ip == '3':
                ip = '192.168.0.200'
                break
            else:
                ip = input('请正确输入1 或者 2 或者 3：')
        return ip

    @staticmethod
    def choose_db_or_col(name, name_num, d_or_c):
        """
        选择数据库或者集合
        :return:
        """
        while True:
            try:
                int(name_num)
            except ValueError:
                name_num = input('请输入正确的%s序号：' % d_or_c)
            if 0 <= int(name_num) <= len(name):
                break
            else:
                name_num = input('请输入正确的%s序号：' % d_or_c)
        for j, e_name in enumerate(name, 1):
            if name_num == str(j):
                choose_name = e_name
                return choose_name

    @staticmethod
    def pr_datas(names):
        """
        规范打印数据库/集合名称
        :param names:
        :return:
        """
        for i, each_name in enumerate(names, 1):
            print(str(i) + '. ' + each_name + ' ' * (40 - len(each_name) - len(str(i))), end='')
            if i % 3 == 0:
                print('\n')
        if len(names) < 3:
            print('\n')

    def get_pdf_msg(self):
        """
        将pdf信息存入消息队列
        :return:
        """
        all_urls = self.choose_way()
        for i in all_urls:
            self.q.put(i)

    def choose_way(self):
        """
        选择PDF的url来源，是在txt文件中还是从数据库中读取
        :return:
        """
        print('*' * 100)
        print(
            '注意事项：' + '\n' + '1. 如果选择了txt文件方式下载，txt内容必须为一行一个url格式' + '\n' +
            '2. 如果选择了从数据库中读取PDF信息下载，则数据库中包含两个字段(a. 存放url的字段，b. 该PDF的名字字段(可以没有))，其中存放url的字段不能嵌套，只能有一个url' + '\n' +
            '3. txt文件只有链接，默认取链接最后 "/" 后的值作为PDF名' + '\n' +
            '4. 数据库方式如果没有指定PDF名字，则默认以链接最后一个 / 后的内容作为PDF名')
        print('*' * 100)
        global is_txt, is_insert, new_collection
        the_way = input('请选择PDF的来源(输入1 / 2)：' + '\n' + '1. txt文件       2. 从数据库中读取' + '\n')
        while True:
            if the_way == '1' or the_way == '2':
                break
            else:
                the_way = input('请正确输入1 或者 2：')
        urls = []
        if the_way == '1':
            # txt文件
            is_txt = True
            print('说明：txt文件中必须为一行一个url！')
            txt_name = input('请输入与该程序同级目录下的txt文件名称(如：123.txt)：')
            while True:
                try:
                    with open(txt_name, 'r') as f:
                        urls = f.readlines()
                    break
                except FileNotFoundError:
                    txt_name = input('没有这个txt文件,请重新输入：')
        elif the_way == '2':
            # 数据库
            collection, db = self.client_dbs()  # 连接数据库
            url_field = input('请输入PDF的url字段名：')
            name_field = input('请输入PDF的name字段名(如果没有则不输入)：')
            if not name_field:
                is_txt = True
                for d in collection.find():
                    if d[url_field]:
                        urls.append(d[url_field])
            else:
                is_txt = False
                for d in collection.find():
                    if d[url_field]:
                        msg = {'pdf_url': d[url_field], 'pdf_name': d[name_field]}
                        urls.append(msg)

        # 选择是否要插入数据的数据库
        is_insert = input('是否要将下载的PDF信息写入到新的数据库(输入y / n，不输入或输入其他默认不写入)：')
        if is_insert == 'y':
            print('选择插入数据库，字段包含(pdf_url, md5_name, relative_path)')
            ip = self.choose_ip()

            client = pymongo.MongoClient('mongodb://%s:27017' % ip)
            db_name = input('请输入要存入的数据库名(可新建)：')
            db = client[db_name]
            if ip == '192.168.0.123':
                db.authenticate(name='用户名', password='密码', source='admin')
            new_col = input('请输入要插入的集合名称：')
            new_collection = db[new_col]
        return urls

    @staticmethod
    def pdf_name_md5(pdf):
        """
        将pdf名字转为md5
        :param pdf:
        :return:
        """
        md = hashlib.md5()
        md.update(pdf.encode('utf-8'))
        pdf_name = md.hexdigest()
        return pdf_name


if __name__ == '__main__':
    download = DownloadPDF()
    download.run()