@@ -34,7 +34,7 @@ def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db
34
34
self .db = redis .StrictRedis (
35
35
host = host , port = port , password = password , db = db , decode_responses = True , ** kwargs )
36
36
37
- def add (self , proxy : Proxy , score = PROXY_SCORE_INIT ) -> int :
37
+ def add (self , proxy : Proxy , score = PROXY_SCORE_INIT , redis_key = REDIS_KEY ) -> int :
38
38
"""
39
39
add proxy and set it to init score
40
40
:param proxy: proxy, ip:port, like 8.8.8.8:88
@@ -44,12 +44,12 @@ def add(self, proxy: Proxy, score=PROXY_SCORE_INIT) -> int:
44
44
if not is_valid_proxy (f'{ proxy .host } :{ proxy .port } ' ):
45
45
logger .info (f'invalid proxy { proxy } , throw it' )
46
46
return
47
- if not self .exists (proxy ):
47
+ if not self .exists (proxy , redis_key ):
48
48
if IS_REDIS_VERSION_2 :
49
- return self .db .zadd (REDIS_KEY , score , proxy .string ())
50
- return self .db .zadd (REDIS_KEY , {proxy .string (): score })
49
+ return self .db .zadd (redis_key , score , proxy .string ())
50
+ return self .db .zadd (redis_key , {proxy .string (): score })
51
51
52
- def random (self ) -> Proxy :
52
+ def random (self , redis_key = REDIS_KEY , proxy_score_min = PROXY_SCORE_MIN , proxy_score_max = PROXY_SCORE_MAX ) -> Proxy :
53
53
"""
54
54
get random proxy
55
55
firstly try to get proxy with max score
@@ -59,74 +59,74 @@ def random(self) -> Proxy:
59
59
"""
60
60
# try to get proxy with max score
61
61
proxies = self .db .zrangebyscore (
62
- REDIS_KEY , PROXY_SCORE_MAX , PROXY_SCORE_MAX )
62
+ redis_key , proxy_score_max , proxy_score_max )
63
63
if len (proxies ):
64
64
return convert_proxy_or_proxies (choice (proxies ))
65
65
# else get proxy by rank
66
66
proxies = self .db .zrevrange (
67
- REDIS_KEY , PROXY_SCORE_MIN , PROXY_SCORE_MAX )
67
+ redis_key , proxy_score_min , proxy_score_max )
68
68
if len (proxies ):
69
69
return convert_proxy_or_proxies (choice (proxies ))
70
70
# else raise error
71
71
raise PoolEmptyException
72
72
73
- def decrease (self , proxy : Proxy ) -> int :
73
+ def decrease (self , proxy : Proxy , redis_key = REDIS_KEY , proxy_score_min = PROXY_SCORE_MIN ) -> int :
74
74
"""
75
75
decrease score of proxy, if small than PROXY_SCORE_MIN, delete it
76
76
:param proxy: proxy
77
77
:return: new score
78
78
"""
79
79
if IS_REDIS_VERSION_2 :
80
- self .db .zincrby (REDIS_KEY , proxy .string (), - 1 )
80
+ self .db .zincrby (redis_key , proxy .string (), - 1 )
81
81
else :
82
- self .db .zincrby (REDIS_KEY , - 1 , proxy .string ())
83
- score = self .db .zscore (REDIS_KEY , proxy .string ())
82
+ self .db .zincrby (redis_key , - 1 , proxy .string ())
83
+ score = self .db .zscore (redis_key , proxy .string ())
84
84
logger .info (f'{ proxy .string ()} score decrease 1, current { score } ' )
85
- if score <= PROXY_SCORE_MIN :
85
+ if score <= proxy_score_min :
86
86
logger .info (f'{ proxy .string ()} current score { score } , remove' )
87
- self .db .zrem (REDIS_KEY , proxy .string ())
87
+ self .db .zrem (redis_key , proxy .string ())
88
88
89
- def exists (self , proxy : Proxy ) -> bool :
89
+ def exists (self , proxy : Proxy , redis_key = REDIS_KEY ) -> bool :
90
90
"""
91
91
if proxy exists
92
92
:param proxy: proxy
93
93
:return: if exists, bool
94
94
"""
95
- return not self .db .zscore (REDIS_KEY , proxy .string ()) is None
95
+ return not self .db .zscore (redis_key , proxy .string ()) is None
96
96
97
- def max (self , proxy : Proxy ) -> int :
97
+ def max (self , proxy : Proxy , redis_key = REDIS_KEY , proxy_score_max = PROXY_SCORE_MAX ) -> int :
98
98
"""
99
99
set proxy to max score
100
100
:param proxy: proxy
101
101
:return: new score
102
102
"""
103
- logger .info (f'{ proxy .string ()} is valid, set to { PROXY_SCORE_MAX } ' )
103
+ logger .info (f'{ proxy .string ()} is valid, set to { proxy_score_max } ' )
104
104
if IS_REDIS_VERSION_2 :
105
- return self .db .zadd (REDIS_KEY , PROXY_SCORE_MAX , proxy .string ())
106
- return self .db .zadd (REDIS_KEY , {proxy .string (): PROXY_SCORE_MAX })
105
+ return self .db .zadd (redis_key , proxy_score_max , proxy .string ())
106
+ return self .db .zadd (redis_key , {proxy .string (): proxy_score_max })
107
107
108
- def count (self ) -> int :
108
+ def count (self , redis_key = REDIS_KEY ) -> int :
109
109
"""
110
110
get count of proxies
111
111
:return: count, int
112
112
"""
113
- return self .db .zcard (REDIS_KEY )
113
+ return self .db .zcard (redis_key )
114
114
115
- def all (self ) -> List [Proxy ]:
115
+ def all (self , redis_key = REDIS_KEY , proxy_score_min = PROXY_SCORE_MIN , proxy_score_max = PROXY_SCORE_MAX ) -> List [Proxy ]:
116
116
"""
117
117
get all proxies
118
118
:return: list of proxies
119
119
"""
120
- return convert_proxy_or_proxies (self .db .zrangebyscore (REDIS_KEY , PROXY_SCORE_MIN , PROXY_SCORE_MAX ))
120
+ return convert_proxy_or_proxies (self .db .zrangebyscore (redis_key , proxy_score_min , proxy_score_max ))
121
121
122
- def batch (self , cursor , count ) -> List [Proxy ]:
122
+ def batch (self , cursor , count , redis_key = REDIS_KEY ) -> List [Proxy ]:
123
123
"""
124
124
get batch of proxies
125
125
:param cursor: scan cursor
126
126
:param count: scan count
127
127
:return: list of proxies
128
128
"""
129
- cursor , proxies = self .db .zscan (REDIS_KEY , cursor , count = count )
129
+ cursor , proxies = self .db .zscan (redis_key , cursor , count = count )
130
130
return cursor , convert_proxy_or_proxies ([i [0 ] for i in proxies ])
131
131
132
132
0 commit comments