From 50dba866fc57238c04ee2af1f4acde9dfb620c75 Mon Sep 17 00:00:00 2001
From: Alexander Scharinger <rng.dynamics@gmail.com>
Date: Thu, 16 Jun 2022 10:59:25 +0200
Subject: [PATCH 1/4] test: itertools.compress

---
 Lib/test/test_itertools.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 23ede28c8eb..cd5b49d16d3 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -500,8 +500,6 @@ def test_combinatorics(self):
                 self.assertEqual(comb, list(filter(set(perm).__contains__, cwr)))     # comb: cwr that is a perm
                 self.assertEqual(comb, sorted(set(cwr) & set(perm)))            # comb: both a cwr and a perm
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_compress(self):
         self.assertEqual(list(compress(data='ABCDEF', selectors=[1,0,1,0,1,1])), list('ACEF'))
         self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))

From 81d2df806a4f9ff8fbb4fa2037659c36512bfd38 Mon Sep 17 00:00:00 2001
From: Alexander Scharinger <rng.dynamics@gmail.com>
Date: Thu, 16 Jun 2022 11:22:36 +0200
Subject: [PATCH 2/4] fix: itertools.compress keyword arguments

---
 vm/src/stdlib/itertools.rs | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
index 26f2710920f..3d6b046d2a4 100644
--- a/vm/src/stdlib/itertools.rs
+++ b/vm/src/stdlib/itertools.rs
@@ -116,15 +116,15 @@ mod decl {
     #[derive(Debug, PyPayload)]
     struct PyItertoolsCompress {
         data: PyIter,
-        selector: PyIter,
+        selectors: PyIter,
     }
 
     #[derive(FromArgs)]
     struct CompressNewArgs {
-        #[pyarg(positional)]
+        #[pyarg(any)]
         data: PyIter,
-        #[pyarg(positional)]
-        selector: PyIter,
+        #[pyarg(any)]
+        selectors: PyIter,
     }
 
     impl Constructor for PyItertoolsCompress {
@@ -132,10 +132,10 @@ mod decl {
 
         fn py_new(
             cls: PyTypeRef,
-            Self::Args { data, selector }: Self::Args,
+            Self::Args { data, selectors }: Self::Args,
             vm: &VirtualMachine,
         ) -> PyResult {
-            PyItertoolsCompress { data, selector }
+            PyItertoolsCompress { data, selectors }
                 .into_ref_with_type(vm, cls)
                 .map(Into::into)
         }
@@ -148,7 +148,7 @@ mod decl {
     impl IterNext for PyItertoolsCompress {
         fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
             loop {
-                let sel_obj = match zelf.selector.next(vm)? {
+                let sel_obj = match zelf.selectors.next(vm)? {
                     PyIterReturn::Return(obj) => obj,
                     PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
                 };

From dbdbdc887f26f1fd1aef130afed8f47d0a2ef34e Mon Sep 17 00:00:00 2001
From: Alexander Scharinger <rng.dynamics@gmail.com>
Date: Thu, 16 Jun 2022 11:36:31 +0200
Subject: [PATCH 3/4] feat: itertools.compress.__reduce__

---
 vm/src/stdlib/itertools.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
index 3d6b046d2a4..9669f8ba053 100644
--- a/vm/src/stdlib/itertools.rs
+++ b/vm/src/stdlib/itertools.rs
@@ -142,7 +142,12 @@ mod decl {
     }
 
     #[pyimpl(with(IterNext, Constructor))]
-    impl PyItertoolsCompress {}
+    impl PyItertoolsCompress {
+        #[pymethod(magic)]
+        fn reduce(zelf: PyRef<Self>) -> (PyTypeRef, (PyIter, PyIter)) {
+            (zelf.class().clone(), (zelf.data.clone(), zelf.selectors.clone()))
+        }
+    }
 
     impl IterNextIterable for PyItertoolsCompress {}
     impl IterNext for PyItertoolsCompress {

From e4267d928071d804adfb89ec5dec1a7db4457628 Mon Sep 17 00:00:00 2001
From: Alexander Scharinger <rng.dynamics@gmail.com>
Date: Thu, 16 Jun 2022 11:44:23 +0200
Subject: [PATCH 4/4] rustfmt

---
 vm/src/stdlib/itertools.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs
index 9669f8ba053..fc7e7434e49 100644
--- a/vm/src/stdlib/itertools.rs
+++ b/vm/src/stdlib/itertools.rs
@@ -145,7 +145,10 @@ mod decl {
     impl PyItertoolsCompress {
         #[pymethod(magic)]
         fn reduce(zelf: PyRef<Self>) -> (PyTypeRef, (PyIter, PyIter)) {
-            (zelf.class().clone(), (zelf.data.clone(), zelf.selectors.clone()))
+            (
+                zelf.class().clone(),
+                (zelf.data.clone(), zelf.selectors.clone()),
+            )
         }
     }