Add chunk utility
This commit is contained in:
parent
69d7f64ead
commit
4c1c95e4ff
|
@ -30,3 +30,49 @@ def round_to(value: float, multiples: float) -> float:
|
|||
if value is None:
|
||||
return None
|
||||
return round(value / multiples) * multiples
|
||||
|
||||
|
||||
def chunk_list(lst, n):
|
||||
for s in range(0, len(lst), n):
|
||||
yield lst[s : s + n]
|
||||
|
||||
|
||||
class chunk:
|
||||
def __init__(self, iterable, n):
|
||||
self.iterable = iterable
|
||||
self.n = n
|
||||
|
||||
def __iter__(self):
|
||||
if isinstance(self.iterable, list):
|
||||
yield from chunk_list(self.iterable, self.n)
|
||||
return
|
||||
|
||||
it = iter(self.iterable)
|
||||
while True:
|
||||
current = []
|
||||
try:
|
||||
for _ in range(self.n):
|
||||
current.append(next(it))
|
||||
yield current
|
||||
except StopIteration:
|
||||
if current:
|
||||
yield current
|
||||
break
|
||||
|
||||
async def __aiter__(self):
|
||||
if hasattr(self.iterable, "__iter__"):
|
||||
for item in self:
|
||||
yield item
|
||||
return
|
||||
|
||||
it = self.iterable.__aiter__()
|
||||
while True:
|
||||
current = []
|
||||
try:
|
||||
for _ in range(self.n):
|
||||
current.append(await it.__anext__())
|
||||
yield current
|
||||
except StopAsyncIteration:
|
||||
if len(current):
|
||||
yield current
|
||||
break
|
||||
|
|
Loading…
Reference in a new issue