File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -400,7 +400,7 @@ def _message_generator(self):
400
400
401
401
tp = self ._next_partition_records .topic_partition
402
402
403
- for msg in self ._next_partition_records .take (sys . maxint ):
403
+ for msg in self ._next_partition_records .take ():
404
404
405
405
# Because we are in a generator, it is possible for
406
406
# subscription state to change between yield calls
@@ -881,9 +881,11 @@ def __len__(self):
881
881
def discard (self ):
882
882
self .messages = None
883
883
884
- def take (self , n ):
884
+ def take (self , n = None ):
885
885
if not len (self ):
886
886
return []
887
+ if n is None or n > len (self ):
888
+ n = len (self )
887
889
next_idx = self .message_idx + n
888
890
res = self .messages [self .message_idx :next_idx ]
889
891
self .message_idx = next_idx
You can’t perform that action at this time.
0 commit comments