@@ -34,7 +34,7 @@ def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db
3434 self .db = redis .StrictRedis (
3535 host = host , port = port , password = password , db = db , decode_responses = True , ** kwargs )
3636
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 :
3838 """
3939 add proxy and set it to init score
4040 :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:
4444 if not is_valid_proxy (f'{ proxy .host } :{ proxy .port } ' ):
4545 logger .info (f'invalid proxy { proxy } , throw it' )
4646 return
47- if not self .exists (proxy ):
47+ if not self .exists (proxy , redis_key ):
4848 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 })
5151
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 :
5353 """
5454 get random proxy
5555 firstly try to get proxy with max score
@@ -59,74 +59,74 @@ def random(self) -> Proxy:
5959 """
6060 # try to get proxy with max score
6161 proxies = self .db .zrangebyscore (
62- REDIS_KEY , PROXY_SCORE_MAX , PROXY_SCORE_MAX )
62+ redis_key , proxy_score_max , proxy_score_max )
6363 if len (proxies ):
6464 return convert_proxy_or_proxies (choice (proxies ))
6565 # else get proxy by rank
6666 proxies = self .db .zrevrange (
67- REDIS_KEY , PROXY_SCORE_MIN , PROXY_SCORE_MAX )
67+ redis_key , proxy_score_min , proxy_score_max )
6868 if len (proxies ):
6969 return convert_proxy_or_proxies (choice (proxies ))
7070 # else raise error
7171 raise PoolEmptyException
7272
73- def decrease (self , proxy : Proxy ) -> int :
73+ def decrease (self , proxy : Proxy , redis_key = REDIS_KEY , proxy_score_min = PROXY_SCORE_MIN ) -> int :
7474 """
7575 decrease score of proxy, if small than PROXY_SCORE_MIN, delete it
7676 :param proxy: proxy
7777 :return: new score
7878 """
7979 if IS_REDIS_VERSION_2 :
80- self .db .zincrby (REDIS_KEY , proxy .string (), - 1 )
80+ self .db .zincrby (redis_key , proxy .string (), - 1 )
8181 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 ())
8484 logger .info (f'{ proxy .string ()} score decrease 1, current { score } ' )
85- if score <= PROXY_SCORE_MIN :
85+ if score <= proxy_score_min :
8686 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 ())
8888
89- def exists (self , proxy : Proxy ) -> bool :
89+ def exists (self , proxy : Proxy , redis_key = REDIS_KEY ) -> bool :
9090 """
9191 if proxy exists
9292 :param proxy: proxy
9393 :return: if exists, bool
9494 """
95- return not self .db .zscore (REDIS_KEY , proxy .string ()) is None
95+ return not self .db .zscore (redis_key , proxy .string ()) is None
9696
97- def max (self , proxy : Proxy ) -> int :
97+ def max (self , proxy : Proxy , redis_key = REDIS_KEY , proxy_score_max = PROXY_SCORE_MAX ) -> int :
9898 """
9999 set proxy to max score
100100 :param proxy: proxy
101101 :return: new score
102102 """
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 } ' )
104104 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 })
107107
108- def count (self ) -> int :
108+ def count (self , redis_key = REDIS_KEY ) -> int :
109109 """
110110 get count of proxies
111111 :return: count, int
112112 """
113- return self .db .zcard (REDIS_KEY )
113+ return self .db .zcard (redis_key )
114114
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 ]:
116116 """
117117 get all proxies
118118 :return: list of proxies
119119 """
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 ))
121121
122- def batch (self , cursor , count ) -> List [Proxy ]:
122+ def batch (self , cursor , count , redis_key = REDIS_KEY ) -> List [Proxy ]:
123123 """
124124 get batch of proxies
125125 :param cursor: scan cursor
126126 :param count: scan count
127127 :return: list of proxies
128128 """
129- cursor , proxies = self .db .zscan (REDIS_KEY , cursor , count = count )
129+ cursor , proxies = self .db .zscan (redis_key , cursor , count = count )
130130 return cursor , convert_proxy_or_proxies ([i [0 ] for i in proxies ])
131131
132132
0 commit comments