@@ -57,15 +57,15 @@ def __init__(self,
5757 # Dictionary from script to sha ''Script''
5858 self .shas = dict ()
5959
60- #### Connection Functions ### #
60+ # Connection Functions #
6161
6262 def echo (self , msg ):
6363 return self ._encode (msg )
6464
6565 def ping (self ):
6666 return b'PONG'
6767
68- #### Transactions Functions ### #
68+ # Transactions Functions #
6969
7070 def lock (self , key , timeout = 0 , sleep = 0 ):
7171 """Emulate lock."""
@@ -101,7 +101,7 @@ def execute(self):
101101 in this mock, so this is a no-op."""
102102 pass
103103
104- #### Keys Functions ### #
104+ # Keys Functions #
105105
106106 def type (self , key ):
107107 key = self ._encode (key )
@@ -247,7 +247,7 @@ def _rename(self, old_key, new_key, nx=False):
247247 return True
248248 return False
249249
250- #### String Functions ### #
250+ # String Functions #
251251
252252 def get (self , key ):
253253 key = self ._encode (key )
@@ -413,7 +413,46 @@ def incr(self, key, amount=1):
413413
414414 incrby = incr
415415
416- #### Hash Functions ####
416+ def setbit (self , key , offset , value ):
417+ """
418+ Set the bit at ``offset`` in ``key`` to ``value``.
419+ """
420+ key = self ._encode (key )
421+ index , bits , mask = self ._get_bits_and_offset (key , offset )
422+
423+ if index >= len (bits ):
424+ bits .extend (b"\x00 " * (index + 1 - len (bits )))
425+
426+ prev_val = 1 if (bits [index ] & mask ) else 0
427+
428+ if value :
429+ bits [index ] |= mask
430+ else :
431+ bits [index ] &= ~ mask
432+
433+ self .redis [key ] = bytes (bits )
434+
435+ return prev_val
436+
437+ def getbit (self , key , offset ):
438+ """
439+ Returns the bit value at ``offset`` in ``key``.
440+ """
441+ key = self ._encode (key )
442+ index , bits , mask = self ._get_bits_and_offset (key , offset )
443+
444+ if index >= len (bits ):
445+ return 0
446+
447+ return 1 if (bits [index ] & mask ) else 0
448+
449+ def _get_bits_and_offset (self , key , offset ):
450+ bits = bytearray (self .redis .get (key , b"" ))
451+ index , position = divmod (offset , 8 )
452+ mask = 128 >> position
453+ return index , bits , mask
454+
455+ # Hash Functions #
417456
418457 def hexists (self , hashkey , attribute ):
419458 """Emulate hexists."""
@@ -518,7 +557,7 @@ def hvals(self, hashkey):
518557 redis_hash = self ._get_hash (hashkey , 'HVALS' )
519558 return redis_hash .values ()
520559
521- #### List Functions ### #
560+ # List Functions #
522561
523562 def lrange (self , key , start , stop ):
524563 """Emulate lrange."""
@@ -726,7 +765,7 @@ def sort(self, name,
726765 return []
727766
728767 by = self ._encode (by ) if by is not None else by
729- # always organize the items as tuples of the value from the list itself and the value to sort by
768+ # always organize the items as tuples of the value from the list and the sort key
730769 if by and b'*' in by :
731770 items = [(i , self .get (by .replace (b'*' , self ._encode (i )))) for i in items ]
732771 elif by in [None , b'nosort' ]:
@@ -782,7 +821,7 @@ def sort(self, name,
782821 else :
783822 return results
784823
785- #### SCAN COMMANDS ### #
824+ # SCAN COMMANDS #
786825
787826 def _common_scan (self , values_function , cursor = '0' , match = None , count = 10 , key = None ):
788827 """
@@ -820,6 +859,14 @@ def value_function():
820859 return sorted (self .redis .keys ()) # sorted list for consistent order
821860 return self ._common_scan (value_function , cursor = cursor , match = match , count = count )
822861
862+ def scan_iter (self , match = None , count = 10 ):
863+ """Emulate scan_iter."""
864+ cursor = '0'
865+ while cursor != 0 :
866+ cursor , data = self .scan (cursor = cursor , match = match , count = count )
867+ for item in data :
868+ yield item
869+
823870 def sscan (self , name , cursor = '0' , match = None , count = 10 ):
824871 """Emulate sscan."""
825872 def value_function ():
@@ -828,13 +875,31 @@ def value_function():
828875 return members
829876 return self ._common_scan (value_function , cursor = cursor , match = match , count = count )
830877
878+ def sscan_iter (self , name , match = None , count = 10 ):
879+ """Emulate sscan_iter."""
880+ cursor = '0'
881+ while cursor != 0 :
882+ cursor , data = self .sscan (name , cursor = cursor ,
883+ match = match , count = count )
884+ for item in data :
885+ yield item
886+
831887 def zscan (self , name , cursor = '0' , match = None , count = 10 ):
832888 """Emulate zscan."""
833889 def value_function ():
834890 values = self .zrange (name , 0 , - 1 , withscores = True )
835891 values .sort (key = lambda x : x [1 ]) # sort for consistent order
836892 return values
837- return self ._common_scan (value_function , cursor = cursor , match = match , count = count , key = lambda v : v [0 ])
893+ return self ._common_scan (value_function , cursor = cursor , match = match , count = count , key = lambda v : v [0 ]) # noqa
894+
895+ def zscan_iter (self , name , match = None , count = 10 ):
896+ """Emulate zscan_iter."""
897+ cursor = '0'
898+ while cursor != 0 :
899+ cursor , data = self .zscan (name , cursor = cursor , match = match ,
900+ count = count )
901+ for item in data :
902+ yield item
838903
839904 def hscan (self , name , cursor = '0' , match = None , count = 10 ):
840905 """Emulate hscan."""
@@ -843,11 +908,20 @@ def value_function():
843908 values = list (values .items ()) # list of tuples for sorting and matching
844909 values .sort (key = lambda x : x [0 ]) # sort for consistent order
845910 return values
846- scanned = self ._common_scan (value_function , cursor = cursor , match = match , count = count , key = lambda v : v [0 ])
911+ scanned = self ._common_scan (value_function , cursor = cursor , match = match , count = count , key = lambda v : v [0 ]) # noqa
847912 scanned [1 ] = dict (scanned [1 ]) # from list of tuples back to dict
848913 return scanned
849914
850- #### SET COMMANDS ####
915+ def hscan_iter (self , name , match = None , count = 10 ):
916+ """Emulate hscan_iter."""
917+ cursor = '0'
918+ while cursor != 0 :
919+ cursor , data = self .hscan (name , cursor = cursor ,
920+ match = match , count = count )
921+ for item in data .items ():
922+ yield item
923+
924+ # SET COMMANDS #
851925
852926 def sadd (self , key , * values ):
853927 """Emulate sadd."""
@@ -960,7 +1034,7 @@ def sunionstore(self, dest, keys, *args):
9601034 self .redis [self ._encode (dest )] = result
9611035 return len (result )
9621036
963- #### SORTED SET COMMANDS ### #
1037+ # SORTED SET COMMANDS #
9641038
9651039 def zadd (self , name , * args , ** kwargs ):
9661040 zset = self ._get_zset (name , "ZADD" , create = True )
@@ -981,21 +1055,21 @@ def zadd(self, name, *args, **kwargs):
9811055 # kwargs
9821056 pieces .extend (kwargs .items ())
9831057
984- insert_count = lambda member , score : 1 if zset .insert (self ._encode (member ), float (score )) else 0
1058+ insert_count = lambda member , score : 1 if zset .insert (self ._encode (member ), float (score )) else 0 # noqa
9851059 return sum ((insert_count (member , score ) for member , score in pieces ))
9861060
9871061 def zcard (self , name ):
9881062 zset = self ._get_zset (name , "ZCARD" )
9891063
9901064 return len (zset ) if zset is not None else 0
9911065
992- def zcount (self , name , min_ , max_ ):
1066+ def zcount (self , name , min , max ):
9931067 zset = self ._get_zset (name , "ZCOUNT" )
9941068
9951069 if not zset :
9961070 return 0
9971071
998- return len (zset .scorerange (float (min_ ), float (max_ )))
1072+ return len (zset .scorerange (float (min ), float (max )))
9991073
10001074 def zincrby (self , name , value , amount = 1 ):
10011075 zset = self ._get_zset (name , "ZINCRBY" , create = True )
@@ -1041,7 +1115,7 @@ def zrange(self, name, start, end, desc=False, withscores=False,
10411115 func = self ._range_func (withscores , score_cast_func )
10421116 return [func (item ) for item in zset .range (start , end , desc )]
10431117
1044- def zrangebyscore (self , name , min_ , max_ , start = None , num = None ,
1118+ def zrangebyscore (self , name , min , max , start = None , num = None ,
10451119 withscores = False , score_cast_func = float ):
10461120 if (start is None ) ^ (num is None ):
10471121 raise RedisError ('`start` and `num` must both be specified' )
@@ -1052,9 +1126,9 @@ def zrangebyscore(self, name, min_, max_, start=None, num=None,
10521126 return []
10531127
10541128 func = self ._range_func (withscores , score_cast_func )
1055- include_start , min_ = self ._score_inclusive (min_ )
1056- include_end , max_ = self ._score_inclusive (max_ )
1057- scorerange = zset .scorerange (min_ , max_ , start_inclusive = include_start , end_inclusive = include_end )
1129+ include_start , min = self ._score_inclusive (min )
1130+ include_end , max = self ._score_inclusive (max )
1131+ scorerange = zset .scorerange (min , max , start_inclusive = include_start , end_inclusive = include_end ) # noqa
10581132 if start is not None and num is not None :
10591133 start , num = self ._translate_limit (len (scorerange ), int (start ), int (num ))
10601134 scorerange = scorerange [start :start + num ]
@@ -1085,23 +1159,23 @@ def zremrangebyrank(self, name, start, end):
10851159
10861160 start , end = self ._translate_range (len (zset ), start , end )
10871161 count_removals = lambda score , member : 1 if zset .remove (member ) else 0
1088- removal_count = sum ((count_removals (score , member ) for score , member in zset .range (start , end )))
1162+ removal_count = sum ((count_removals (score , member ) for score , member in zset .range (start , end ))) # noqa
10891163 if removal_count > 0 and len (zset ) == 0 :
10901164 self .delete (name )
10911165 return removal_count
10921166
1093- def zremrangebyscore (self , name , min_ , max_ ):
1167+ def zremrangebyscore (self , name , min , max ):
10941168 zset = self ._get_zset (name , "ZREMRANGEBYSCORE" )
10951169
10961170 if not zset :
10971171 return 0
10981172
10991173 count_removals = lambda score , member : 1 if zset .remove (member ) else 0
1100- include_start , min_ = self ._score_inclusive (min_ )
1101- include_end , max_ = self ._score_inclusive (max_ )
1174+ include_start , min = self ._score_inclusive (min )
1175+ include_end , max = self ._score_inclusive (max )
11021176
11031177 removal_count = sum ((count_removals (score , member )
1104- for score , member in zset .scorerange (min_ , max_ ,
1178+ for score , member in zset .scorerange (min , max ,
11051179 start_inclusive = include_start ,
11061180 end_inclusive = include_end )))
11071181 if removal_count > 0 and len (zset ) == 0 :
@@ -1113,7 +1187,7 @@ def zrevrange(self, name, start, end, withscores=False,
11131187 return self .zrange (name , start , end ,
11141188 desc = True , withscores = withscores , score_cast_func = score_cast_func )
11151189
1116- def zrevrangebyscore (self , name , max_ , min_ , start = None , num = None ,
1190+ def zrevrangebyscore (self , name , max , min , start = None , num = None ,
11171191 withscores = False , score_cast_func = float ):
11181192
11191193 if (start is None ) ^ (num is None ):
@@ -1124,11 +1198,12 @@ def zrevrangebyscore(self, name, max_, min_, start=None, num=None,
11241198 return []
11251199
11261200 func = self ._range_func (withscores , score_cast_func )
1127- include_start , min_ = self ._score_inclusive (min_ )
1128- include_end , max_ = self ._score_inclusive (max_ )
1201+ include_start , min = self ._score_inclusive (min )
1202+ include_end , max = self ._score_inclusive (max )
11291203
1130- scorerange = [x for x in reversed (zset .scorerange (float (min_ ), float (max_ ),
1131- start_inclusive = include_start , end_inclusive = include_end ))]
1204+ scorerange = [x for x in reversed (zset .scorerange (float (min ), float (max ),
1205+ start_inclusive = include_start ,
1206+ end_inclusive = include_end ))]
11321207 if start is not None and num is not None :
11331208 start , num = self ._translate_limit (len (scorerange ), int (start ), int (num ))
11341209 scorerange = scorerange [start :start + num ]
@@ -1170,7 +1245,7 @@ def zunionstore(self, dest, keys, aggregate=None):
11701245 self .redis [self ._encode (dest )] = union
11711246 return len (union )
11721247
1173- #### Script Commands ### #
1248+ # Script Commands #
11741249
11751250 def eval (self , script , numkeys , * keys_and_args ):
11761251 """Emulate eval"""
@@ -1279,12 +1354,12 @@ def _normalize_command_response(self, command, response):
12791354
12801355 return response
12811356
1282- #### PubSub commands ### #
1357+ # PubSub commands #
12831358
12841359 def publish (self , channel , message ):
12851360 self .pubsub [channel ].append (message )
12861361
1287- #### Internal ### #
1362+ # Internal #
12881363
12891364 def _get_list (self , key , operation , create = False ):
12901365 """
@@ -1308,7 +1383,7 @@ def _get_zset(self, name, operation, create=False):
13081383 """
13091384 Get (and maybe create) a sorted set by name.
13101385 """
1311- return self ._get_by_type (name , operation , create , b'zset' , SortedSet (), return_default = False )
1386+ return self ._get_by_type (name , operation , create , b'zset' , SortedSet (), return_default = False ) # noqa
13121387
13131388 def _get_by_type (self , key , operation , create , type_ , default , return_default = True ):
13141389 """
@@ -1348,7 +1423,7 @@ def _range_func(self, withscores, score_cast_func):
13481423 Return a suitable function from (score, member)
13491424 """
13501425 if withscores :
1351- return lambda score_member : (score_member [1 ], score_cast_func (self ._encode (score_member [0 ])))
1426+ return lambda score_member : (score_member [1 ], score_cast_func (self ._encode (score_member [0 ]))) # noqa
13521427 else :
13531428 return lambda score_member : score_member [1 ]
13541429
0 commit comments