Skip to content

Commit 1ccce90

Browse files
nonzerofloatyouknowone
authored andcommitted
Add itertools.chain.__setstate__
1 parent bfa595d commit 1ccce90

2 files changed

Lines changed: 31 additions & 2 deletions

File tree

‎Lib/test/test_itertools.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,8 +195,7 @@ def test_chain_reducible(self):
195195
self.assertRaises(TypeError, list, oper(chain(2, 3)))
196196
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
197197
self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
198-
# TODO: RUSTPYTHON
199-
@unittest.expectedFailure
198+
200199
def test_chain_setstate(self):
201200
self.assertRaises(TypeError, chain().__setstate__, ())
202201
self.assertRaises(TypeError, chain().__setstate__, [])
@@ -209,6 +208,7 @@ def test_chain_setstate(self):
209208
it = chain()
210209
it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
211210
self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
211+
212212
# TODO: RUSTPYTHON
213213
@unittest.expectedFailure
214214
def test_combinations(self):

‎vm/src/stdlib/itertools.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,36 @@ mod decl {
5858
fn class_getitem(cls: PyTypeRef, args: PyObjectRef, vm: &VirtualMachine) -> PyGenericAlias {
5959
PyGenericAlias::new(cls, args, vm)
6060
}
61+
62+
#[pymethod(magic)]
63+
fn setstate(&self, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> {
64+
let args = state.as_slice();
65+
if args.is_empty() {
66+
let msg = String::from("function takes at leat 1 arguments (0 given)");
67+
return Err(vm.new_type_error(msg));
68+
}
69+
if args.len() > 2 {
70+
let msg = format!("function takes at most 2 arguments ({} given)", args.len());
71+
return Err(vm.new_type_error(msg));
72+
}
73+
let source = &args[0];
74+
let active = args.get(1);
75+
if !PyIter::check(source.as_ref())
76+
|| !active.map_or(true, |active| PyIter::check(active.as_ref()))
77+
{
78+
return Err(vm.new_type_error(String::from("Arguments must be iterators.")));
79+
}
80+
*self.iterables.write() = source.try_to_value(vm)?;
81+
self.cur_idx.store(0);
82+
*self.cached_iter.write() = if let Some(active) = active {
83+
Some(active.clone().get_iter(vm)?)
84+
} else {
85+
None
86+
};
87+
Ok(())
88+
}
6189
}
90+
6291
impl IterNextIterable for PyItertoolsChain {}
6392
impl IterNext for PyItertoolsChain {
6493
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {

0 commit comments

Comments
 (0)