-
Notifications
You must be signed in to change notification settings - Fork 542
Expand file tree
/
Copy pathwebdriver_pool.py
More file actions
115 lines (97 loc) · 3.21 KB
/
webdriver_pool.py
File metadata and controls
115 lines (97 loc) · 3.21 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# -*- coding: utf-8 -*-
"""
Created on 2021/3/18 4:59 下午
---------
@summary:
---------
@author: Boris
@email: boris_liu@foxmail.com
"""
import queue
import threading
from feapder.utils.log import log
from feapder.utils.tools import Singleton
from feapder.utils.webdriver.selenium_driver import SeleniumDriver
@Singleton
class WebDriverPool:
def __init__(
self, pool_size=5, driver_cls=SeleniumDriver, thread_safe=False, **kwargs
):
"""
Args:
pool_size: driver池的大小
driver: 驱动类型
thread_safe: 是否线程安全
是则每个线程拥有一个driver，pool_size无效，driver数量为线程数
否则每个线程从池中获取driver
**kwargs:
"""
self.pool_size = pool_size
self.driver_cls = driver_cls
self.thread_safe = thread_safe
self.kwargs = kwargs
self.queue = queue.Queue(maxsize=pool_size)
self.lock = threading.RLock()
self.driver_count = 0
self.ctx = threading.local()
@property
def driver(self):
if not hasattr(self.ctx, "driver"):
self.ctx.driver = None
return self.ctx.driver
@driver.setter
def driver(self, driver):
self.ctx.driver = driver
@property
def is_full(self):
return self.driver_count >= self.pool_size
def create_driver(self, user_agent: str = None, proxy: str = None):
kwargs = self.kwargs.copy()
if user_agent:
kwargs["user_agent"] = user_agent
if proxy:
kwargs["proxy"] = proxy
return self.driver_cls(**kwargs)
def get(self, user_agent: str = None, proxy: str = None):
"""
获取webdriver
当webdriver为新实例时会使用 user_agen, proxy, cookie参数来创建
Args:
user_agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/73.0.3683.103 Safari/537.36
proxy: xxx.xxx.xxx.xxx
Returns:
"""
if not self.is_full and not self.thread_safe:
with self.lock:
if not self.is_full:
driver = self.create_driver(user_agent, proxy)
self.queue.put(driver)
self.driver_count += 1
elif self.thread_safe:
if not self.driver:
driver = self.create_driver(user_agent, proxy)
self.driver = driver
self.driver_count += 1
if self.thread_safe:
driver = self.driver
else:
driver = self.queue.get()
return driver
def put(self, driver):
if not self.thread_safe:
self.queue.put(driver)
def remove(self, driver):
if self.thread_safe:
if self.driver:
self.driver.quit()
self.driver = None
else:
driver.quit()
self.driver_count -= 1
def close(self):
if self.thread_safe:
log.info("暂不支持关闭需线程安全的driver")
while not self.queue.empty():
driver = self.queue.get()
driver.quit()
self.driver_count -= 1