from pypy.interpreter.baseobjspace import W_Root from pypy.interpreter.error import OperationError, oefmt from pypy.interpreter.typedef import TypeDef, make_weakref_descr from pypy.interpreter.gateway import interp2app, unwrap_spec, WrappedDefault from pypy.objspace.std.util import generic_alias_class_getitem, builtinclass_new_args_check from rpython.rlib import jit from rpython.rtyper.lltypesystem import rffi from pypy.module.__builtin__.functional import W_Filter, build_iterators_from_args def W_Twoarg__new__(space, w_subtype, W_Base, name, __args__): args_w = __args__.arguments_w w_type = space.gettypeobject(W_Base.typedef) w_init = space.newtext("__init__") if (space.is_w(w_subtype, w_type) or space.is_w(space.getattr(w_subtype, w_init), space.getattr(w_type, w_init))): if __args__.keyword_names_w: raise oefmt(space.w_TypeError, "%s() takes no keyword arguments", name) length = len(args_w) if args_w else 0 if length != 2: raise oefmt(space.w_TypeError, "%s() expected 2 arguments, got %d", name, length) return args_w class W_Count(W_Root): def __init__(self, space, w_firstval, w_step): self.space = space self.w_c = w_firstval self.w_step = w_step def iter_w(self): return self def next_w(self): w_c = self.w_c self.w_c = self.space.add(w_c, self.w_step) return w_c def single_argument(self): space = self.space return (space.isinstance_w(self.w_step, space.w_int) and space.eq_w(self.w_step, space.newint(1))) def repr_w(self): space = self.space cls_name = space.type(self).getname(space) c = space.text_w(space.repr(self.w_c)) if self.single_argument(): s = '%s(%s)' % (cls_name, c) else: step = space.text_w(space.repr(self.w_step)) s = '%s(%s, %s)' % (cls_name, c, step) return self.space.newtext(s) def reduce_w(self): space = self.space if self.single_argument(): args_w = [self.w_c] else: args_w = [self.w_c, self.w_step] return space.newtuple2(space.gettypefor(W_Count), space.newtuple(args_w)) def check_number(space, w_obj): if (space.lookup(w_obj, '__int__') is None and space.lookup(w_obj, '__float__') is None and not space.isinstance_w(w_obj, space.w_complex) ): raise oefmt(space.w_TypeError, "expected a number") @unwrap_spec(w_start=WrappedDefault(0), w_step=WrappedDefault(1)) def W_Count___new__(space, w_subtype, w_start, w_step): check_number(space, w_start) check_number(space, w_step) r = space.allocate_instance(W_Count, w_subtype) r.__init__(space, w_start, w_step) return r W_Count.typedef = TypeDef( 'itertools.count', __new__ = interp2app(W_Count___new__), __iter__ = interp2app(W_Count.iter_w), __next__ = interp2app(W_Count.next_w), __reduce__ = interp2app(W_Count.reduce_w), __repr__ = interp2app(W_Count.repr_w), __doc__ = """Make an iterator that returns evenly spaced values starting with n. If not specified n defaults to zero. Often used as an argument to imap() to generate consecutive data points. Also, used with izip() to add sequence numbers. Equivalent to: def count(start=0, step=1): n = start while True: yield n n += step """) class W_Repeat(W_Root): def __init__(self, space, w_obj, w_times): self.space = space self.w_obj = w_obj if w_times is None: self.counting = False self.count = 0 else: self.counting = True self.count = max(self.space.int_w(w_times), 0) def next_w(self): if self.counting: if self.count <= 0: raise OperationError(self.space.w_StopIteration, self.space.w_None) self.count -= 1 return self.w_obj def iter_w(self): return self def length_w(self): if not self.counting: return self.space.w_NotImplemented return self.space.newint(self.count) def repr_w(self): space = self.space cls_name = space.type(self).getname(space) objrepr = self.space.text_w(space.repr(self.w_obj)) if self.counting: s = '%s(%s, %d)' % (cls_name, objrepr, self.count) else: s = '%s(%s)' % (cls_name, objrepr) return self.space.newtext(s) def descr_reduce(self): space = self.space if self.counting: args_w = [self.w_obj, space.newint(self.count)] else: args_w = [self.w_obj] return space.newtuple([space.gettypefor(W_Repeat), space.newtuple(args_w)]) def W_Repeat___new__(space, w_subtype, w_object, w_times=None): r = space.allocate_instance(W_Repeat, w_subtype) r.__init__(space, w_object, w_times) return r W_Repeat.typedef = TypeDef( 'itertools.repeat', __new__ = interp2app(W_Repeat___new__), __iter__ = interp2app(W_Repeat.iter_w), __length_hint__ = interp2app(W_Repeat.length_w), __next__ = interp2app(W_Repeat.next_w), __repr__ = interp2app(W_Repeat.repr_w), __reduce__ = interp2app(W_Repeat.descr_reduce), __doc__ = """Make an iterator that returns object over and over again. Runs indefinitely unless the times argument is specified. Used as argument to imap() for invariant parameters to the called function. Equivalent to : def repeat(object, times=None): if times is None: while True: yield object else: for i in xrange(times): yield object """) class W_TakeWhile(W_Root): def __init__(self, space, w_predicate, w_iterable): self.space = space self.w_predicate = w_predicate self.w_iterable = space.iter(w_iterable) self.stopped = False def iter_w(self): return self def next_w(self): if self.stopped: raise OperationError(self.space.w_StopIteration, self.space.w_None) w_obj = self.space.next(self.w_iterable) # may raise a w_StopIteration w_bool = self.space.call_function(self.w_predicate, w_obj) if not self.space.is_true(w_bool): self.stopped = True raise OperationError(self.space.w_StopIteration, self.space.w_None) return w_obj def descr_reduce(self, space): return space.newtuple([ space.type(self), space.newtuple([self.w_predicate, self.w_iterable]), space.newbool(self.stopped) ]) def descr_setstate(self, space, w_state): self.stopped = space.bool_w(w_state) def W_TakeWhile___new__(space, w_subtype, __args__): args_w = W_Twoarg__new__(space, w_subtype, W_TakeWhile, "takewhile", __args__) r = space.allocate_instance(W_TakeWhile, w_subtype) r.__init__(space, args_w[0], args_w[1]) return r W_TakeWhile.typedef = TypeDef( 'itertools.takewhile', __new__ = interp2app(W_TakeWhile___new__), __iter__ = interp2app(W_TakeWhile.iter_w), __next__ = interp2app(W_TakeWhile.next_w), __reduce__ = interp2app(W_TakeWhile.descr_reduce), __setstate__ = interp2app(W_TakeWhile.descr_setstate), __doc__ = """Make an iterator that returns elements from the iterable as long as the predicate is true. Equivalent to : def takewhile(predicate, iterable): for x in iterable: if predicate(x): yield x else: break """) class W_DropWhile(W_Root): def __init__(self, space, w_predicate, w_iterable): self.space = space self.w_predicate = w_predicate self.w_iterable = space.iter(w_iterable) self.started = False def iter_w(self): return self def next_w(self): if self.started: w_obj = self.space.next(self.w_iterable) # may raise w_StopIter else: while True: w_obj = self.space.next(self.w_iterable) # may raise w_StopIter w_bool = self.space.call_function(self.w_predicate, w_obj) if not self.space.is_true(w_bool): self.started = True break return w_obj def descr_reduce(self, space): return space.newtuple([ space.type(self), space.newtuple([self.w_predicate, self.w_iterable]), space.newbool(self.started) ]) def descr_setstate(self, space, w_state): self.started = space.bool_w(w_state) def W_DropWhile___new__(space, w_subtype, __args__): args_w = W_Twoarg__new__(space, w_subtype, W_DropWhile, "dropwhile", __args__) r = space.allocate_instance(W_DropWhile, w_subtype) r.__init__(space, args_w[0], args_w[1]) return r W_DropWhile.typedef = TypeDef( 'itertools.dropwhile', __new__ = interp2app(W_DropWhile___new__), __iter__ = interp2app(W_DropWhile.iter_w), __next__ = interp2app(W_DropWhile.next_w), __reduce__ = interp2app(W_DropWhile.descr_reduce), __setstate__ = interp2app(W_DropWhile.descr_setstate), __doc__ = """Make an iterator that drops elements from the iterable as long as the predicate is true; afterwards, returns every element. Note, the iterator does not produce any output until the predicate is true, so it may have a lengthy start-up time. Equivalent to : def dropwhile(predicate, iterable): iterable = iter(iterable) for x in iterable: if not predicate(x): yield x break for x in iterable: yield x """) class W_FilterFalse(W_Filter): reverse = True def descr_reduce(self, space): args_w = [space.w_None if self.w_predicate is None else self.w_predicate, self.w_iterable] return space.newtuple([space.type(self), space.newtuple(args_w)]) def W_FilterFalse___new__(space, w_subtype, __args__): args_w = W_Twoarg__new__(space, w_subtype, W_FilterFalse, 'filterfalse', __args__) r = space.allocate_instance(W_FilterFalse, w_subtype) r.__init__(space, args_w[0], args_w[1]) return r W_FilterFalse.typedef = TypeDef( 'itertools.filterfalse', __new__ = interp2app(W_FilterFalse___new__), __iter__ = interp2app(W_FilterFalse.iter_w), __next__ = interp2app(W_FilterFalse.next_w), __reduce__ = interp2app(W_FilterFalse.descr_reduce), __doc__ = """Make an iterator that filters elements from iterable returning only those for which the predicate is False. If predicate is None, return the items that are false. Equivalent to : def filterfalse(predicate, iterable): if predicate is None: predicate = bool for x in iterable: if not predicate(x): yield x """) def get_printable_location(greenkey): return "islice_ignore_items [%s]" % (greenkey.iterator_greenkey_printable(), ) islice_ignore_items_driver = jit.JitDriver(name='islice_ignore_items', greens=['greenkey'], reds=['w_islice', 'w_iterator'], get_printable_location=get_printable_location) class W_ISlice(W_Root): def __init__(self, space, w_iterable, w_startstop, args_w): self.iterable = space.iter(w_iterable) self.space = space num_args = len(args_w) if num_args == 0: start = 0 w_stop = w_startstop elif num_args <= 2: if space.is_w(w_startstop, space.w_None): start = 0 else: start = self.arg_int_w(w_startstop, 0, "Indicies for islice() must be None or non-negative integers") w_stop = args_w[0] else: raise oefmt(space.w_TypeError, "islice() takes at most 4 arguments (%d given)", num_args) if space.is_w(w_stop, space.w_None): stop = -1 else: stop = self.arg_int_w(w_stop, 0, "Stop argument must be a non-negative integer or None.") stop = max(start, stop) # for obscure CPython compatibility if num_args == 2: w_step = args_w[1] if space.is_w(w_step, space.w_None): step = 1 else: step = self.arg_int_w(w_step, 1, "Step for islice() must be a positive integer or None") else: step = 1 self.count = rffi.cast(rffi.UNSIGNED, 0) self.next = rffi.cast(rffi.UNSIGNED, start) self.stop = stop self.step = rffi.cast(rffi.UNSIGNED, step) def arg_int_w(self, w_obj, minimum, errormsg): space = self.space try: result = space.int_w(space.index(w_obj)) except OperationError as e: if e.async(space): raise result = -1 if result < minimum: raise OperationError(space.w_ValueError, space.newtext(errormsg)) return result def iter_w(self): return self def next_w(self): if self.iterable is None: raise OperationError(self.space.w_StopIteration, self.space.w_None) if self.count != self.next: self._ignore_items() stop = self.stop if 0 <= stop <= self.count: self.iterable = None raise OperationError(self.space.w_StopIteration, self.space.w_None) try: item = self.space.next(self.iterable) except OperationError as e: if e.match(self.space, self.space.w_StopIteration): self.iterable = None raise self.count += 1 oldnext = self.next new_next = self.next + self.step if new_next < oldnext or (stop >=0 and new_next > rffi.cast(rffi.UNSIGNED,stop)): self.next = rffi.cast(rffi.UNSIGNED, stop) else: self.next = new_next return item def _ignore_items(self): w_iterator = self.iterable greenkey = self.space.iterator_greenkey(w_iterator) while True: islice_ignore_items_driver.jit_merge_point( greenkey=greenkey, w_islice=self, w_iterator=w_iterator) if self.count >= self.next: break try: self.space.next(w_iterator) except OperationError as e: if e.match(self.space, self.space.w_StopIteration): self.iterable = None raise self.count += 1 def descr_reduce(self, space): if self.iterable is None: return space.newtuple([ space.type(self), space.newtuple([space.iter(space.newlist([])), space.newint(0)]), space.newint(0), ]) stop = self.stop if stop == -1: w_stop = space.w_None else: w_stop = space.newint(stop) return space.newtuple([ space.type(self), space.newtuple([self.iterable, space.newint(self.next), w_stop, space.newint(self.step)]), space.newint(self.count), ]) def descr_setstate(self, space, w_state): self.count = rffi.cast(rffi.UNSIGNED, space.int_w(w_state)) def W_ISlice___new__(space, w_subtype, w_iterable, w_startstop, __args__): args_w = __args__.arguments_w w_islice = space.gettypeobject(W_ISlice.typedef) w_init = space.newtext("__init__") if (space.is_w(w_subtype, w_islice) or space.is_w(space.getattr(w_subtype, w_init), space.getattr(w_islice, w_init))): if __args__.keyword_names_w: raise oefmt(space.w_TypeError, "islice() takes no keyword arguments") r = space.allocate_instance(W_ISlice, w_subtype) r.__init__(space, w_iterable, w_startstop, args_w) return r W_ISlice.typedef = TypeDef( 'itertools.islice', __new__ = interp2app(W_ISlice___new__), __iter__ = interp2app(W_ISlice.iter_w), __next__ = interp2app(W_ISlice.next_w), __reduce__ = interp2app(W_ISlice.descr_reduce), __setstate__ = interp2app(W_ISlice.descr_setstate), __doc__ = """Make an iterator that returns selected elements from the iterable. If start is non-zero, then elements from the iterable are skipped until start is reached. Afterward, elements are returned consecutively unless step is set higher than one which results in items being skipped. If stop is None, then iteration continues until the iterator is exhausted, if at all; otherwise, it stops at the specified position. Unlike regular slicing, islice() does not support negative values for start, stop, or step. Can be used to extract related fields from data where the internal structure has been flattened (for example, a multi-line report may list a name field on every third line). """) class W_Chain(W_Root): def __init__(self, space, w_iterables): self.space = space self.w_iterables = w_iterables self.w_it = None def iter_w(self): return self def _advance(self): if self.w_iterables is None: raise OperationError(self.space.w_StopIteration, self.space.w_None) self.w_it = self.space.iter(self.space.next(self.w_iterables)) def next_w(self): if not self.w_it: try: self._advance() except OperationError as e: raise e try: return self.space.next(self.w_it) except OperationError as e: return self._handle_error(e) def _handle_error(self, e): while True: if not e.match(self.space, self.space.w_StopIteration): raise e try: self._advance() # may raise StopIteration itself except OperationError as e: self.w_iterables = None raise e try: return self.space.next(self.w_it) except OperationError as e: pass # loop back to the start of _handle_error(e) def descr_reduce(self, space): if self.w_iterables is not None: if self.w_it is not None: inner_contents = [self.w_iterables, self.w_it] else: inner_contents = [self.w_iterables] result_w = [space.type(self), space.newtuple([]), space.newtuple(inner_contents)] else: result_w = [space.type(self), space.newtuple([])] return space.newtuple(result_w) def descr_setstate(self, space, w_state): state = space.unpackiterable(w_state) num_args = len(state) if num_args < 1: raise oefmt(space.w_TypeError, "function takes at least 1 argument (%d given)", num_args) elif num_args == 1: self.w_iterables = state[0] elif num_args == 2: self.w_iterables, self.w_it = state else: raise oefmt(space.w_TypeError, "function takes at most 2 arguments (%d given)", num_args) def W_Chain___new__(space, w_subtype, __args__): args_w = __args__.arguments_w w_chain = space.gettypeobject(W_Chain.typedef) w_init = space.newtext("__init__") if (space.is_w(w_subtype, w_chain) or space.is_w(space.getattr(w_subtype, w_init), space.getattr(w_chain, w_init))): if __args__.keyword_names_w: raise oefmt(space.w_TypeError, "chain() takes no keyword arguments") r = space.allocate_instance(W_Chain, w_subtype) w_args = space.newtuple(args_w) r.__init__(space, space.iter(w_args)) return r def chain_from_iterable(space, w_cls, w_arg): """chain.from_iterable(iterable) --> chain object Alternate chain() constructor taking a single iterable argument that evaluates lazily.""" r = space.allocate_instance(W_Chain, w_cls) r.__init__(space, space.iter(w_arg)) return r W_Chain.typedef = TypeDef( 'itertools.chain', __new__ = interp2app(W_Chain___new__), __iter__ = interp2app(W_Chain.iter_w), __next__ = interp2app(W_Chain.next_w), __reduce__ = interp2app(W_Chain.descr_reduce), __setstate__ = interp2app(W_Chain.descr_setstate), from_iterable = interp2app(chain_from_iterable, as_classmethod=True), __class_getitem__ = interp2app( generic_alias_class_getitem, as_classmethod=True), __doc__ = """Make an iterator that returns elements from the first iterable until it is exhausted, then proceeds to the next iterable, until all of the iterables are exhausted. Used for treating consecutive sequences as a single sequence. Equivalent to : def chain(*iterables): for it in iterables: for element in it: yield element """) class W_ZipLongest(W_Root): _immutable_fields_ = ["w_fillvalue", "iterators"] def __init__(self, space, w_fun, args_w): self.space = space self.w_fun = w_fun self.iterators_w = build_iterators_from_args(space, args_w, "zip_longest") def iter_w(self): return self def _fetch(self, index): w_iter = self.iterators_w[index] if w_iter is not None: space = self.space try: return space.next(w_iter) except OperationError as e: if not e.match(space, space.w_StopIteration): raise self.active -= 1 if self.active <= 0: # It was the last active iterator raise self.iterators_w[index] = None return self.w_fillvalue def next_w(self): # common case: 2 arguments if len(self.iterators_w) == 2: return self.space.newtuple2(self._fetch(0), self._fetch(1)) else: objects = self._get_objects() return self.space.newtuple(objects) def _get_objects(self): # the loop is out of the way of the JIT nb = len(self.iterators_w) if nb == 0: raise OperationError(self.space.w_StopIteration, self.space.w_None) return [self._fetch(index) for index in range(nb)] def descr_reduce(self, space): result_w = [space.type(self)] if self.iterators_w is not None: iterators = [iterator if iterator is not None else space.newtuple([]) for iterator in self.iterators_w] iterators = space.newtuple(iterators) else: iterators = space.newtuple([]) result_w = [space.type(self), iterators, self.w_fillvalue] return space.newtuple(result_w) def descr_setstate(self, space, w_state): self.w_fillvalue = w_state def iterator_greenkey(self, space): # XXX in theory we should tupleize the greenkeys of all the # sub-iterators, but much more work if len(self.iterators_w) > 0: return space.iterator_greenkey(self.iterators_w[0]) return None def W_ZipLongest___new__(space, w_subtype, __args__): arguments_w, kwds_w = __args__.unpack() w_fillvalue = space.w_None if kwds_w: if "fillvalue" in kwds_w: w_fillvalue = kwds_w["fillvalue"] del kwds_w["fillvalue"] if kwds_w: raise oefmt(space.w_TypeError, "zip_longest() got unexpected keyword argument(s)") self = space.allocate_instance(W_ZipLongest, w_subtype) self.__init__(space, space.w_None, arguments_w) self.w_fillvalue = w_fillvalue self.active = len(self.iterators_w) return self W_ZipLongest.typedef = TypeDef( 'itertools.zip_longest', __new__ = interp2app(W_ZipLongest___new__), __iter__ = interp2app(W_ZipLongest.iter_w), __next__ = interp2app(W_ZipLongest.next_w), __reduce__ = interp2app(W_ZipLongest.descr_reduce), __setstate__ = interp2app(W_ZipLongest.descr_setstate), __doc__ = """Return a zip_longest object whose .next() method returns a tuple where the i-th element comes from the i-th iterable argument. The .next() method continues until the longest iterable in the argument sequence is exhausted and then it raises StopIteration. When the shorter iterables are exhausted, the fillvalue is substituted in their place. The fillvalue defaults to None or can be specified by a keyword argument. """) class W_Cycle(W_Root): def __init__(self, space, w_iterable): self.space = space self.saved_w = [] self.w_iterable = space.iter(w_iterable) self.index = 0 # 0 during the first iteration; > 0 afterwards def iter_w(self): return self def next_w(self): if self.index > 0: if not self.saved_w: raise OperationError(self.space.w_StopIteration, self.space.w_None) try: w_obj = self.saved_w[self.index] except IndexError: self.index = 1 w_obj = self.saved_w[0] else: self.index += 1 else: try: w_obj = self.space.next(self.w_iterable) except OperationError as e: if e.match(self.space, self.space.w_StopIteration): self.index = 1 if not self.saved_w: raise w_obj = self.saved_w[0] else: raise else: self.saved_w.append(w_obj) return w_obj def descr_reduce(self, space): # reduces differently than CPython 3.5. Unsure if it is a # problem. To be on the safe side, keep three arguments for # __setstate__; CPython takes two. return space.newtuple([ space.type(self), space.newtuple([self.w_iterable]), space.newtuple([ space.newlist(self.saved_w), space.newint(self.index), # space.newbool(self.index > 0), ]), ]) def descr_setstate(self, space, w_state): state_w = space.unpackiterable(w_state, 2) w_saved = state_w[0] self.saved_w = space.unpackiterable(w_saved) w_index = state_w[1] self.index = space.int_w(w_index) # w_exhausted ignored # w_exhausted = state_w[2] def W_Cycle___new__(space, w_subtype, w_iterable, __posonly__, __args__): w_type = space.gettypeobject(W_Cycle.typedef) builtinclass_new_args_check(space, "cycle", w_type, w_subtype, __args__) r = space.allocate_instance(W_Cycle, w_subtype) r.__init__(space, w_iterable) return r W_Cycle.typedef = TypeDef( 'itertools.cycle', __new__ = interp2app(W_Cycle___new__), __iter__ = interp2app(W_Cycle.iter_w), __next__ = interp2app(W_Cycle.next_w), __reduce__ = interp2app(W_Cycle.descr_reduce), __setstate__ = interp2app(W_Cycle.descr_setstate), __doc__ = """Make an iterator returning elements from the iterable and saving a copy of each. When the iterable is exhausted, return elements from the saved copy. Repeats indefinitely. Equivalent to : def cycle(iterable): saved = [] for element in iterable: yield element saved.append(element) while saved: for element in saved: yield element """) class W_StarMap(W_Root): def __init__(self, space, w_fun, w_iterable): self.space = space self.w_fun = w_fun self.w_iterable = self.space.iter(w_iterable) def iter_w(self): return self def next_w(self): w_obj = self.space.next(self.w_iterable) return self.space.call(self.w_fun, w_obj) def descr_reduce(self): return self.space.newtuple([self.space.gettypefor(W_StarMap), self.space.newtuple([ self.w_fun, self.w_iterable]) ]) def W_StarMap___new__(space, w_subtype, __args__): args_w = W_Twoarg__new__(space, w_subtype, W_StarMap, "starmap", __args__) r = space.allocate_instance(W_StarMap, w_subtype) r.__init__(space, args_w[0], args_w[1]) return r W_StarMap.typedef = TypeDef( 'itertools.starmap', __new__ = interp2app(W_StarMap___new__), __iter__ = interp2app(W_StarMap.iter_w), __next__ = interp2app(W_StarMap.next_w), __reduce__ = interp2app(W_StarMap.descr_reduce), __doc__ = """Make an iterator that computes the function using arguments tuples obtained from the iterable. Used instead of imap() when argument parameters are already grouped in tuples from a single iterable (the data has been ``pre-zipped''). The difference between imap() and starmap() parallels the distinction between function(a,b) and function(*c). Equivalent to : def starmap(function, iterable): iterable = iter(iterable) while True: yield function(*iterable.next()) """) @unwrap_spec(n=int) def tee(space, w_iterable, n=2): """Return n independent iterators from a single iterable. Note : once tee() has made a split, the original iterable should not be used anywhere else; otherwise, the iterable could get advanced without the tee objects being informed. Note : this member of the toolkit may require significant auxiliary storage (depending on how much temporary data needs to be stored). In general, if one iterator is going to use most or all of the data before the other iterator, it is faster to use list() instead of tee() If iter(iterable) has no method __copy__(), this is equivalent to: def tee(iterable, n=2): def gen(next, data={}, cnt=[0]): for i in count(): if i == cnt[0]: item = data[i] = next() cnt[0] += 1 else: item = data[i] # data.pop(i) if it's the last one yield item it = iter(iterable) return tuple([gen(it.next) for i in range(n)]) If iter(iterable) has a __copy__ method, though, we just return a tuple t = (iterable, t[0].__copy__(), t[1].__copy__(), ...). """ if n < 0: raise oefmt(space.w_ValueError, "n must be >= 0") w_iterator = space.iter(w_iterable) if space.findattr(w_iterator, space.newtext("__copy__")) is not None: # In this case, we don't instantiate any W_TeeIterable. # We just rely on doing repeated __copy__(). This case # includes the situation where w_iterable is already # a W_TeeIterable itself. iterators_w = [w_iterator] * n for i in range(1, n): iterators_w[i] = space.call_method(w_iterator, "__copy__") else: w_chained_list = W_TeeChainedListNode(space) iterators_w = [W_TeeIterable(space, w_iterator, w_chained_list) for x in range(n)] return space.newtuple(iterators_w) class W_TeeChainedListNode(W_Root): def __init__(self, space): self.w_next = None self.w_obj = None self.running = False def reduce_w(self, space): list_w = [] node = self while node is not None and node.w_obj is not None: list_w.append(node.w_obj) node = node.w_next if not list_w: return space.newtuple([space.type(self), space.newtuple([])]) return space.newtuple( [space.type(self), space.newtuple([]), space.newtuple([space.newlist(list_w)]) ]) def descr_setstate(self, space, w_state): state = space.unpackiterable(w_state) if len(state) != 1: raise oefmt(space.w_ValueError, "invalid arguments") obj_list_w = space.unpackiterable(state[0]) node = self for w_obj in obj_list_w: node.w_obj = w_obj node.w_next = W_TeeChainedListNode(space) node = node.w_next def W_TeeChainedListNode___new__(space, w_subtype): r = space.allocate_instance(W_TeeChainedListNode, w_subtype) r.__init__(space) return r W_TeeChainedListNode.typedef = TypeDef( 'itertools._tee_dataobject', __new__ = interp2app(W_TeeChainedListNode___new__), __weakref__ = make_weakref_descr(W_TeeChainedListNode), __reduce__ = interp2app(W_TeeChainedListNode.reduce_w), __setstate__ = interp2app(W_TeeChainedListNode.descr_setstate) ) W_TeeChainedListNode.typedef.acceptable_as_base_class = False class W_TeeIterable(W_Root): def __init__(self, space, w_iterator, w_chained_list=None): self.space = space self.w_iterator = w_iterator self.w_chained_list = w_chained_list def iter_w(self): return self def next_w(self): w_chained_list = self.w_chained_list if w_chained_list is None: raise OperationError(self.space.w_StopIteration, self.space.w_None) if w_chained_list.running: raise oefmt(self.space.w_RuntimeError, "cannot re-enter the tee iterator") w_obj = w_chained_list.w_obj if w_obj is None: w_chained_list.running = True try: w_obj = self.space.next(self.w_iterator) w_chained_list.running = False except OperationError as e: if e.match(self.space, self.space.w_StopIteration): self.w_chained_list = None w_chained_list.running = False raise w_chained_list.w_next = W_TeeChainedListNode(self.space) w_chained_list.w_obj = w_obj self.w_chained_list = w_chained_list.w_next return w_obj def copy_w(self): space = self.space tee_iter = W_TeeIterable(space, self.w_iterator, self.w_chained_list) return tee_iter def reduce_w(self): return self.space.newtuple([self.space.gettypefor(W_TeeIterable), self.space.newtuple([self.space.newtuple([])]), self.space.newtuple([ self.w_iterator, self.w_chained_list]) ]) def setstate_w(self, w_state): state = self.space.unpackiterable(w_state) num_args = len(state) if num_args != 2: raise oefmt(self.space.w_TypeError, "function takes exactly 2 arguments (%d given)", num_args) w_iterator, w_chained_list = state if not isinstance(w_chained_list, W_TeeChainedListNode): raise oefmt(self.space.w_TypeError, "must be itertools._tee_dataobject, not %s", self.space.type(w_chained_list).name) self.w_iterator = w_iterator self.w_chained_list = w_chained_list def W_TeeIterable___new__(space, w_subtype, w_iterable): if isinstance(w_iterable, W_TeeIterable): myiter = space.interp_w(W_TeeIterable, w_iterable) w_iterator = myiter.w_iterator w_chained_list = myiter.w_chained_list else: w_iterator = space.iter(w_iterable) w_chained_list = W_TeeChainedListNode(space) return W_TeeIterable(space, w_iterator, w_chained_list) W_TeeIterable.typedef = TypeDef( 'itertools._tee', __new__ = interp2app(W_TeeIterable___new__), __iter__ = interp2app(W_TeeIterable.iter_w), __next__ = interp2app(W_TeeIterable.next_w), __copy__ = interp2app(W_TeeIterable.copy_w), __weakref__ = make_weakref_descr(W_TeeIterable), __reduce__ = interp2app(W_TeeIterable.reduce_w), __setstate__ = interp2app(W_TeeIterable.setstate_w) ) W_TeeIterable.typedef.acceptable_as_base_class = False class W_GroupBy(W_Root): def __init__(self, space, w_iterable, w_fun): self.space = space self.w_iterator = self.space.iter(w_iterable) if w_fun is None: w_fun = space.w_None self.w_keyfunc = w_fun self.w_tgtkey = None self.w_currkey = None self.w_currvalue = None def iter_w(self): return self def next_w(self): self.w_currgrouper = None self._skip_to_next_iteration_group() w_key = self.w_tgtkey = self.w_currkey w_grouper = W_GroupByIterator(self, w_key) return self.space.newtuple2(w_key, w_grouper) def _skip_to_next_iteration_group(self): space = self.space while True: if self.w_currkey is None: pass elif self.w_tgtkey is None: break else: if not space.eq_w(self.w_tgtkey, self.w_currkey): break w_newvalue = space.next(self.w_iterator) if space.is_w(self.w_keyfunc, space.w_None): w_newkey = w_newvalue else: w_newkey = space.call_function(self.w_keyfunc, w_newvalue) self.w_currkey = w_newkey self.w_currvalue = w_newvalue def descr_reduce(self, space): items_w = [space.type(self), space.newtuple([ self.w_iterator, self.w_keyfunc])] if (self.w_tgtkey is not None and self.w_currkey is not None and self.w_currvalue is not None): items_w = items_w + [ space.newtuple([ self.w_currkey, self.w_currvalue, self.w_tgtkey]) ] return space.newtuple(items_w) def descr_setstate(self, space, w_state): state = space.unpackiterable(w_state) num_args = len(state) if num_args != 3: raise oefmt(space.w_TypeError, "function takes exactly 3 arguments (%d given)", num_args) self.w_currkey, self.w_currvalue, self.w_tgtkey = state def W_GroupBy___new__(space, w_subtype, w_iterable, w_key=None): r = space.allocate_instance(W_GroupBy, w_subtype) r.__init__(space, w_iterable, w_key) return r W_GroupBy.typedef = TypeDef( 'itertools.groupby', __new__ = interp2app(W_GroupBy___new__), __iter__ = interp2app(W_GroupBy.iter_w), __next__ = interp2app(W_GroupBy.next_w), __reduce__ = interp2app(W_GroupBy.descr_reduce), __setstate__ = interp2app(W_GroupBy.descr_setstate), __doc__ = """Make an iterator that returns consecutive keys and groups from the iterable. The key is a function computing a key value for each element. If not specified or is None, key defaults to an identity function and returns the element unchanged. Generally, the iterable needs to already be sorted on the same key function. The returned group is itself an iterator that shares the underlying iterable with groupby(). Because the source is shared, when the groupby object is advanced, the previous group is no longer visible. So, if that data is needed later, it should be stored as a list: groups = [] uniquekeys = [] for k, g in groupby(data, keyfunc): groups.append(list(g)) # Store group iterator as a list uniquekeys.append(k) """) class W_GroupByIterator(W_Root): def __init__(self, groupby, w_tgtkey): self.groupby = groupby self.w_tgtkey = w_tgtkey groupby.w_currgrouper = self def iter_w(self): return self def next_w(self): groupby = self.groupby space = groupby.space if groupby.w_currgrouper is not self: raise OperationError(space.w_StopIteration, space.w_None) if groupby.w_currvalue is None: w_newvalue = space.next(groupby.w_iterator) if space.is_w(groupby.w_keyfunc, space.w_None): w_newkey = w_newvalue else: w_newkey = space.call_function(groupby.w_keyfunc, w_newvalue) #assert groupby.w_currvalue is None # ^^^ check disabled, see http://bugs.python.org/issue30347 groupby.w_currkey = w_newkey groupby.w_currvalue = w_newvalue assert groupby.w_currkey is not None if not space.eq_w(self.w_tgtkey, groupby.w_currkey): raise OperationError(space.w_StopIteration, space.w_None) w_result = groupby.w_currvalue groupby.w_currvalue = None groupby.w_currkey = None return w_result def descr_reduce(self, space): if self.groupby.w_currgrouper is not self: w_callable = space.builtin.get('iter') return space.newtuple([w_callable, space.newtuple([space.newtuple([])])]) return space.newtuple([ space.type(self), space.newtuple([ self.groupby, self.w_tgtkey]), ]) def W_GroupByIterator__new__(space, w_subtype, w_parent, w_tgtkey): r = space.allocate_instance(W_GroupByIterator, w_subtype) groupby = space.interp_w(W_GroupBy, w_parent) r.__init__(groupby, w_tgtkey) return r W_GroupByIterator.typedef = TypeDef( 'itertools._groupby', __new__ = interp2app(W_GroupByIterator__new__), __iter__ = interp2app(W_GroupByIterator.iter_w), __next__ = interp2app(W_GroupByIterator.next_w), __reduce__ = interp2app(W_GroupByIterator.descr_reduce)) W_GroupByIterator.typedef.acceptable_as_base_class = False class W_Compress(W_Root): def __init__(self, space, w_data, w_selectors): self.space = space self.w_data = space.iter(w_data) self.w_selectors = space.iter(w_selectors) def iter_w(self): return self def next_w(self): # No need to check for StopIteration since either w_data # or w_selectors will raise this. The shortest one stops first. while True: w_next_item = self.space.next(self.w_data) w_next_selector = self.space.next(self.w_selectors) if self.space.is_true(w_next_selector): return w_next_item def descr_reduce(self, space): return space.newtuple([ space.type(self), space.newtuple([self.w_data, self.w_selectors]) ]) def W_Compress__new__(space, w_subtype, w_data, w_selectors): r = space.allocate_instance(W_Compress, w_subtype) r.__init__(space, w_data, w_selectors) return r W_Compress.typedef = TypeDef( 'itertools.compress', __new__ = interp2app(W_Compress__new__), __iter__ = interp2app(W_Compress.iter_w), __next__ = interp2app(W_Compress.next_w), __reduce__ = interp2app(W_Compress.descr_reduce), __doc__ = """Make an iterator that filters elements from *data* returning only those that have a corresponding element in *selectors* that evaluates to ``True``. Stops when either the *data* or *selectors* iterables has been exhausted. Equivalent to:: def compress(data, selectors): # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F return (d for d, s in zip(data, selectors) if s) """) class W_Product(W_Root): def __init__(self, space, args_w, w_repeat): self.gears = [ space.unpackiterable(arg_w)[:] for arg_w in args_w ] * space.int_w(w_repeat) # for gear in self.gears: if len(gear) == 0: self.indices = None self.lst = None self.stopped = True break else: self.indices = [0] * len(self.gears) self.lst = None self.stopped = False def _rotate_previous_gears(self): lst = self.lst x = len(self.gears) - 1 lst[x] = self.gears[x][0] self.indices[x] = 0 x -= 1 # the outer loop runs as long as a we have a carry while x >= 0: gear = self.gears[x] index = self.indices[x] + 1 if index < len(gear): # no carry: done lst[x] = gear[index] self.indices[x] = index return lst[x] = gear[0] self.indices[x] = 0 x -= 1 else: self.lst = None self.stopped = True def fill_next_result(self): # the last gear is done here, in a function with no loop, # to allow the JIT to look inside if self.lst is None: self.lst = [None for gear in self.gears] for index, gear in enumerate(self.gears): self.lst[index] = gear[0] return lst = self.lst x = len(self.gears) - 1 if x >= 0: gear = self.gears[x] index = self.indices[x] + 1 if index < len(gear): # no carry: done lst[x] = gear[index] self.indices[x] = index else: self._rotate_previous_gears() else: self.stopped = True def iter_w(self, space): return self def next_w(self, space): if not self.stopped: self.fill_next_result() if self.stopped: raise OperationError(space.w_StopIteration, space.w_None) w_result = space.newtuple(self.lst[:]) return w_result def descr_reduce(self, space): if not self.stopped: gears = [space.newtuple(gear) for gear in self.gears] result_w = [ space.type(self), space.newtuple(gears) ] if self.lst is not None: indices_w = [space.newint(index) for index in self.indices] result_w = result_w + [space.newtuple(indices_w)] else: result_w = [ space.type(self), space.newtuple([space.newtuple([])]) ] return space.newtuple(result_w) def descr_setstate(self, space, w_state): gear_count = len(self.gears) indices_w = space.unpackiterable(w_state) lst = [] for i, gear in enumerate(self.gears): w_index = indices_w[i] index = space.int_w(w_index) gear_size = len(gear) if self.indices is None or gear_size == 0: self.stopped = True return if index < 0: index = 0 if index > gear_size - 1: index = gear_size - 1 self.indices[i] = index lst.append(gear[index]) self.lst = lst def W_Product__new__(space, w_subtype, __args__): arguments_w, kwds_w = __args__.unpack() w_repeat = space.newint(1) if kwds_w: if 'repeat' in kwds_w: w_repeat = kwds_w['repeat'] del kwds_w['repeat'] if kwds_w: raise oefmt(space.w_TypeError, "product() got unexpected keyword argument(s)") r = space.allocate_instance(W_Product, w_subtype) r.__init__(space, arguments_w, w_repeat) return r W_Product.typedef = TypeDef( 'itertools.product', __new__ = interp2app(W_Product__new__), __iter__ = interp2app(W_Product.iter_w), __next__ = interp2app(W_Product.next_w), __reduce__ = interp2app(W_Product.descr_reduce), __setstate__ = interp2app(W_Product.descr_setstate), __doc__ = """ Cartesian product of input iterables. Equivalent to nested for-loops in a generator expression. For example, ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``. The nested loops cycle like an odometer with the rightmost element advancing on every iteration. This pattern creates a lexicographic ordering so that if the input's iterables are sorted, the product tuples are emitted in sorted order. To compute the product of an iterable with itself, specify the number of repetitions with the optional *repeat* keyword argument. For example, ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``. This function is equivalent to the following code, except that the actual implementation does not build up intermediate results in memory:: def product(*args, **kwds): # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111 pools = map(tuple, args) * kwds.get('repeat', 1) result = [[]] for pool in pools: result = [x+[y] for x in result for y in pool] for prod in result: yield tuple(prod) """) class W_Combinations(W_Root): def __init__(self, space, pool_w, indices, r): self.pool_w = pool_w self.indices = indices self.r = r self.last_result_w = None self.stopped = r > len(pool_w) def get_maximum(self, i): return i + len(self.pool_w) - self.r def max_index(self, j): return self.indices[j - 1] + 1 def descr__iter__(self, space): return self def descr_next(self, space): if self.stopped: raise OperationError(space.w_StopIteration, space.w_None) if self.last_result_w is None: # On the first pass, initialize result tuple using the indices result_w = [None] * self.r for i in xrange(self.r): index = self.indices[i] result_w[i] = self.pool_w[index] else: # Copy the previous result result_w = self.last_result_w[:] # Scan indices right-to-left until finding one that is not at its # maximum i = self.r - 1 while i >= 0 and self.indices[i] == self.get_maximum(i): i -= 1 # If i is negative, then the indices are all at their maximum value # and we're done if i < 0: self.stopped = True raise OperationError(space.w_StopIteration, space.w_None) # Increment the current index which we know is not at its maximum. # Then move back to the right setting each index to its lowest # possible value self.indices[i] += 1 for j in xrange(i + 1, self.r): self.indices[j] = self.max_index(j) # Update the result for the new indices starting with i, the # leftmost index that changed for i in xrange(i, self.r): index = self.indices[i] w_elem = self.pool_w[index] result_w[i] = w_elem self.last_result_w = result_w return space.newtuple(result_w) def descr_reduce(self, space): if self.stopped: pool_w = [] else: pool_w = self.pool_w result_w = [ space.type(self), space.newtuple([ space.newtuple(pool_w), space.newint(self.r) ])] if self.last_result_w is not None and not self.stopped: # we must pickle the indices and use them for setstate result_w = result_w + [ space.newtuple([ space.newint(index) for index in self.indices])] return space.newtuple(result_w) def descr_setstate(self, space, w_state): indices_w = space.fixedview(w_state) if len(indices_w) != self.r: raise oefmt(space.w_ValueError, "invalid arguments") for i in range(self.r): index = space.int_w(indices_w[i]) max = self.get_maximum(i) # clamp the index (beware of negative max) if index > max: index = max if index < 0: index = 0 self.indices[i] = index self.last_result_w = [ self.pool_w[self.indices[i]] for i in range(self.r)] @unwrap_spec(r=int) def W_Combinations__new__(space, w_subtype, w_iterable, r): pool_w = space.fixedview(w_iterable) if r < 0: raise oefmt(space.w_ValueError, "r must be non-negative") indices = range(r) res = space.allocate_instance(W_Combinations, w_subtype) res.__init__(space, pool_w, indices, r) return res W_Combinations.typedef = TypeDef("itertools.combinations", __new__ = interp2app(W_Combinations__new__), __iter__ = interp2app(W_Combinations.descr__iter__), __next__ = interp2app(W_Combinations.descr_next), __reduce__ = interp2app(W_Combinations.descr_reduce), __setstate__ = interp2app(W_Combinations.descr_setstate), __doc__ = """\ combinations(iterable, r) --> combinations object Return successive r-length combinations of elements in the iterable. combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)""", ) class W_CombinationsWithReplacement(W_Combinations): def __init__(self, space, pool_w, indices, r): W_Combinations.__init__(self, space, pool_w, indices, r) self.stopped = len(pool_w) == 0 and r > 0 def get_maximum(self, i): return len(self.pool_w) - 1 def max_index(self, j): return self.indices[j - 1] def descr_reduce(self, space): if self.stopped: pool_w = [] else: pool_w = self.pool_w result_w = [ space.type(self), space.newtuple([ space.newtuple(pool_w), space.newint(self.r) ])] if self.last_result_w is not None and not self.stopped: # we must pickle the indices and use them for setstate result_w = result_w + [ space.newtuple([ space.newint(index) for index in self.indices])] return space.newtuple(result_w) def descr_setstate(self, space, w_state): indices_w = space.fixedview(w_state) if len(indices_w) != self.r: raise oefmt(space.w_ValueError, "invalid arguments") for i in range(self.r): index = space.int_w(indices_w[i]) max = self.get_maximum(i) # clamp the index (beware of negative max) if index > max: index = max if index < 0: index = 0 self.indices[i] = index self.last_result_w = [ self.pool_w[self.indices[i]] for i in range(self.r)] @unwrap_spec(r=int) def W_CombinationsWithReplacement__new__(space, w_subtype, w_iterable, r): pool_w = space.fixedview(w_iterable) if r < 0: raise oefmt(space.w_ValueError, "r must be non-negative") indices = [0] * r res = space.allocate_instance(W_CombinationsWithReplacement, w_subtype) res.__init__(space, pool_w, indices, r) return res W_CombinationsWithReplacement.typedef = TypeDef( "itertools.combinations_with_replacement", __new__ = interp2app(W_CombinationsWithReplacement__new__), __iter__ = interp2app(W_CombinationsWithReplacement.descr__iter__), __next__ = interp2app(W_CombinationsWithReplacement.descr_next), __reduce__ = interp2app(W_CombinationsWithReplacement.descr_reduce), __setstate__ = interp2app(W_CombinationsWithReplacement.descr_setstate), __doc__ = """\ combinations_with_replacement(iterable, r) --> combinations_with_replacement object Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats. combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC""", ) class W_Permutations(W_Root): def __init__(self, space, pool_w, r): self.pool_w = pool_w self.r = r n = len(pool_w) n_minus_r = n - r if n_minus_r < 0: self.stopped = self.raised_stop_iteration = True else: self.stopped = self.raised_stop_iteration = False self.indices = range(n) self.cycles = range(n, n_minus_r, -1) self.started = False def descr__iter__(self, space): return self def descr_next(self, space): if self.stopped: self.raised_stop_iteration = True raise OperationError(space.w_StopIteration, space.w_None) r = self.r indices = self.indices w_result = space.newtuple([self.pool_w[indices[i]] for i in range(r)]) cycles = self.cycles i = r - 1 while i >= 0: j = cycles[i] - 1 if j > 0: cycles[i] = j indices[i], indices[-j] = indices[-j], indices[i] return w_result cycles[i] = len(indices) - i n1 = len(indices) - 1 assert n1 >= 0 num = indices[i] for k in range(i, n1): indices[k] = indices[k+1] indices[n1] = num i -= 1 self.stopped = True if self.started: raise OperationError(space.w_StopIteration, space.w_None) else: self.started = True return w_result def descr_reduce(self, space): if self.raised_stop_iteration: pool_w = [] else: pool_w = self.pool_w result_w = [ space.type(self), space.newtuple([ space.newtuple(pool_w), space.newint(self.r) ])] if not self.raised_stop_iteration: # we must pickle the indices and use them for setstate result_w = result_w + [ space.newtuple([ space.newtuple([ space.newint(index) for index in self.indices]), space.newtuple([ space.newint(num) for num in self.cycles]), space.newbool(self.started) ])] return space.newtuple(result_w) def descr_setstate(self, space, w_state): state = space.unpackiterable(w_state) if len(state) == 3: w_indices, w_cycles, w_started = state indices_w = space.unpackiterable(w_indices) cycles_w = space.unpackiterable(w_cycles) self.started = space.bool_w(w_started) else: raise oefmt(space.w_ValueError, "invalid arguments") if len(indices_w) != len(self.pool_w) or len(cycles_w) != self.r: raise oefmt(space.w_ValueError, "inavalid arguments") n = len(self.pool_w) for i in range(n): index = space.int_w(indices_w[i]) if index < 0: index = 0 elif index > n-1: index = n-1 self.indices[i] = index for i in range(self.r): index = space.int_w(cycles_w[i]) if index < 1: index = 1 elif index > n-i: index = n-i self.cycles[i] = index def W_Permutations__new__(space, w_subtype, w_iterable, w_r=None): pool_w = space.fixedview(w_iterable) if space.is_none(w_r): r = len(pool_w) else: r = space.gateway_nonnegint_w(w_r) res = space.allocate_instance(W_Permutations, w_subtype) res.__init__(space, pool_w, r) return res W_Permutations.typedef = TypeDef("itertools.permutations", __new__ = interp2app(W_Permutations__new__), __iter__ = interp2app(W_Permutations.descr__iter__), __next__ = interp2app(W_Permutations.descr_next), __reduce__ = interp2app(W_Permutations.descr_reduce), __setstate__ = interp2app(W_Permutations.descr_setstate), __doc__ = """\ permutations(iterable[, r]) --> permutations object Return successive r-length permutations of elements in the iterable. permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)""", ) class W_Accumulate(W_Root): 'Return series of accumulated sums (or other binary function results).' def __init__(self, space, w_iterable, w_func, w_initial): self.space = space self.w_iterable = w_iterable self.w_func = w_func if not space.is_w(w_func, space.w_None) else None self.w_total = None self.w_initial = w_initial def iter_w(self): return self def next_w(self): space = self.space if not space.is_w(self.w_initial, space.w_None): w_res = self.w_total = self.w_initial self.w_initial = space.w_None return w_res w_value = space.next(self.w_iterable) if self.w_total is None: self.w_total = w_value return w_value if self.w_func is None: self.w_total = space.add(self.w_total, w_value) else: self.w_total = space.call_function(self.w_func, self.w_total, w_value) return self.w_total def reduce_w(self): space = self.space w_func = space.w_None if self.w_func is None else self.w_func if not space.is_w(self.w_initial, space.w_None): w_it = W_Chain(space, space.iter(space.newlist([ space.newtuple([self.w_initial]), self.w_iterable]))) return space.newtuple([space.gettypefor(W_Accumulate), space.newtuple([w_it, w_func]), space.w_None]) if self.w_total is space.w_None: # :-( w_it = W_Chain(space, space.iter(space.newlist([ space.newtuple([self.w_total]), self.w_iterable]))) w_it = space.call_function(space.type(self), w_it, w_func) return space.newtuple([space.gettypefor(W_ISlice), space.newtuple([w_it, space.newint(1), space.w_None])]) w_total = space.w_None if self.w_total is None else self.w_total return space.newtuple([space.gettypefor(W_Accumulate), space.newtuple([self.w_iterable, w_func]), w_total]) def setstate_w(self, space, w_state): self.w_total = w_state if not space.is_w(w_state, space.w_None) else None @unwrap_spec(w_initial=WrappedDefault(None)) def W_Accumulate__new__(space, w_subtype, w_iterable, w_func=None, __kwonly__=None, w_initial=None): r = space.allocate_instance(W_Accumulate, w_subtype) r.__init__(space, space.iter(w_iterable), w_func, w_initial) return r W_Accumulate.typedef = TypeDef("itertools.accumulate", __new__ = interp2app(W_Accumulate__new__), __iter__ = interp2app(W_Accumulate.iter_w), __next__ = interp2app(W_Accumulate.next_w), __reduce__ = interp2app(W_Accumulate.reduce_w), __setstate__ = interp2app(W_Accumulate.setstate_w), __doc__ = """\ "accumulate(iterable) --> accumulate object Return series of accumulated sums.""") class W_Pairwise(W_Root): 'Return series of accumulated sums (or other binary function results).' def __init__(self, space, w_iterator): self.space = space w_prev = None self.w_iterator = w_iterator self.w_prev = None def iter_w(self): return self def next_w(self): space = self.space w_prev = self.w_prev if w_prev is None: w_prev = space.next(self.w_iterator) w_next = space.next(self.w_iterator) w_res = space.newtuple2(w_prev, w_next) self.w_prev = w_next return w_res def W_Pairwise__new__(space, w_subtype, w_iterable, __posonly__=None): r = space.allocate_instance(W_Pairwise, w_subtype) r.__init__(space, space.iter(w_iterable)) return r W_Pairwise.typedef = TypeDef("itertools.pairwise", __new__ = interp2app(W_Pairwise__new__), __iter__ = interp2app(W_Pairwise.iter_w), __next__ = interp2app(W_Pairwise.next_w), __doc__ = """\ Return an iterator of overlapping pairs taken from the input iterator. s -> (s0, s1), (s1, s2), (s2, s3), ... """)