Skip to content

Commit 9a61b6d

Browse files
committed
add two more missing itertools
1 parent 90ba46d commit 9a61b6d

File tree

1 file changed

+67
-0
lines changed

1 file changed

+67
-0
lines changed

graalpython/lib-graalpython/itertools.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,3 +578,70 @@ def __next__(self):
578578
result[i] = elem
579579
self.last_result = result
580580
return tuple(result)
581+
582+
583+
class combinations_with_replacement(combinations):
584+
"""
585+
combinations_with_replacement(iterable, r) --> combinations_with_replacement object
586+
587+
Return successive r-length combinations of elements in the iterable
588+
allowing individual elements to have successive repeats.
589+
combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
590+
"""
591+
def __init__(self, iterable, r):
592+
pool = list(iterable)
593+
if r < 0:
594+
raise ValueError("r must be non-negative")
595+
indices = [0] * r
596+
combinations.__init__(pool, indices, r)
597+
self.stopped = len(pool) == 0 and r > 0
598+
599+
def get_maximum(self, i):
600+
return len(self.pool) - 1
601+
602+
def max_index(self, j):
603+
return self.indices[j - 1]
604+
605+
606+
class zip_longest():
607+
"""
608+
zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object
609+
610+
Return a zip_longest object whose .next() method returns a tuple where
611+
the i-th element comes from the i-th iterable argument. The .next()
612+
method continues until the longest iterable in the argument sequence
613+
is exhausted and then it raises StopIteration. When the shorter iterables
614+
are exhausted, the fillvalue is substituted in their place. The fillvalue
615+
defaults to None or can be specified by a keyword argument.
616+
"""
617+
618+
def __iter__(self):
619+
return self
620+
621+
def _fetch(self, index):
622+
it = self.iterators[index]
623+
if it is not None:
624+
try:
625+
return next(it)
626+
except StopIteration:
627+
self.active -= 1
628+
if self.active <= 0:
629+
# It was the last active iterator
630+
raise
631+
self.iterators[index] = None
632+
return self.fillvalue
633+
634+
def __next__(self):
635+
if self.active <= 0:
636+
raise StopIteration
637+
nb = len(self.iterators)
638+
if nb == 0:
639+
raise StopIteration
640+
return tuple(self._fetch(index) for index in range(nb))
641+
642+
def __new__(subtype, iter1, *args, fillvalue=None):
643+
self = object.__new__(subtype)
644+
self.fillvalue = fillvalue
645+
self.active = len(args) + 1
646+
self.iterators = [iter1] + args
647+
return self

0 commit comments

Comments
 (0)