Author: Amaury Forgeot d'Arc <amaur...@gmail.com>
Branch: py3.3
Changeset: r81996:e4a76a0698fc
Date: 2016-01-27 09:35 +0100
http://bitbucket.org/pypy/pypy/changeset/e4a76a0698fc/
Log: Add pickle support to itertools.combinations (and probably
combinations_with_replacement as well)
diff --git a/pypy/module/itertools/interp_itertools.py
b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -1088,6 +1088,42 @@
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.wrap(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.wrap(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:
+ import pdb;pdb.set_trace()
+ raise OperationError(space.w_ValueError, space.wrap(
+ "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.append(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)
@@ -1095,7 +1131,7 @@
raise OperationError(space.w_ValueError,
space.wrap("r must be non-negative")
)
- indices = range(len(pool_w))
+ indices = range(r)
res = space.allocate_instance(W_Combinations, w_subtype)
res.__init__(space, pool_w, indices, r)
return space.wrap(res)
@@ -1104,6 +1140,8 @@
__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
diff --git a/pypy/module/itertools/test/test_itertools.py
b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -975,6 +975,25 @@
islice = itertools.islice(myiter, 5, 8)
raises(StopIteration, islice.__next__)
+ def test_combinations_pickle(self):
+ from itertools import combinations
+ import pickle
+ for op in (lambda a:a, lambda a:pickle.loads(pickle.dumps(a))):
+ assert list(op(combinations('abc', 32))) == [] # r > n
+ assert list(op(combinations('ABCD', 2))) == [
+ ('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'),
('C','D')]
+ testIntermediate = combinations('ABCD', 2)
+ next(testIntermediate)
+ assert list(op(testIntermediate)) == [
+ ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')]
+
+ assert list(op(combinations(range(4), 3))) == [
+ (0,1,2), (0,1,3), (0,2,3), (1,2,3)]
+ testIntermediate = combinations(range(4), 3)
+ next(testIntermediate)
+ assert list(op(testIntermediate)) == [
+ (0,1,3), (0,2,3), (1,2,3)]
+
class AppTestItertools32:
spaceconfig = dict(usemodules=['itertools'])
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
