From 97fd6585936f6d5f8eb45fdb98f15a281ba1d8d0 Mon Sep 17 00:00:00 2001
From: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
Date: Sun, 12 Apr 2020 06:44:01 +0100
Subject: [PATCH] Add IntoPy<PyObject> for HashSet and BTreeSet (#865)

---
 CHANGELOG.md     |  1 +
 src/types/set.rs | 52 +++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bff307e4..d02a3a57 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ### Added
 * `_PyDict_NewPresized`. [#849](https://github.com/PyO3/pyo3/pull/849)
+* `IntoPy<PyObject>` for `HashSet` and `BTreeSet`. [#864](https://github.com/PyO3/pyo3/pull/864)
 
 ### Fixed
 * `__radd__` and other `__r*__` methods now correctly work with operators. [#839](https://github.com/PyO3/pyo3/pull/839)
diff --git a/src/types/set.rs b/src/types/set.rs
index 6b7aa248..40f8dadd 100644
--- a/src/types/set.rs
+++ b/src/types/set.rs
@@ -4,8 +4,8 @@
 use crate::err::{self, PyErr, PyResult};
 use crate::internal_tricks::Unsendable;
 use crate::{
-    ffi, AsPyPointer, FromPyObject, PyAny, PyNativeType, PyObject, Python, ToBorrowedObject,
-    ToPyObject,
+    ffi, AsPyPointer, FromPy, FromPyObject, IntoPy, PyAny, PyNativeType, PyObject, Python,
+    ToBorrowedObject, ToPyObject,
 };
 use std::cmp;
 use std::collections::{BTreeSet, HashSet};
@@ -184,6 +184,22 @@ where
     }
 }
 
+impl<K, S> FromPy<HashSet<K, S>> for PyObject
+where
+    K: IntoPy<PyObject> + Eq + hash::Hash,
+    S: hash::BuildHasher + Default,
+{
+    fn from_py(src: HashSet<K, S>, py: Python) -> Self {
+        let set = PySet::empty(py).expect("Failed to construct empty set");
+        {
+            for val in src {
+                set.add(val.into_py(py)).expect("Failed to add to set");
+            }
+        }
+        set.into()
+    }
+}
+
 impl<'source, K, S> FromPyObject<'source> for HashSet<K, S>
 where
     K: FromPyObject<'source> + cmp::Eq + hash::Hash,
@@ -195,6 +211,21 @@ where
     }
 }
 
+impl<K> FromPy<BTreeSet<K>> for PyObject
+where
+    K: IntoPy<PyObject> + cmp::Ord + ToPyObject,
+{
+    fn from_py(src: BTreeSet<K>, py: Python) -> Self {
+        let set = PySet::empty(py).expect("Failed to construct empty set");
+        {
+            for val in src {
+                set.add(val.into_py(py)).expect("Failed to add to set");
+            }
+        }
+        set.into()
+    }
+}
+
 impl<'source, K> FromPyObject<'source> for BTreeSet<K>
 where
     K: FromPyObject<'source> + cmp::Ord,
@@ -271,7 +302,7 @@ impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
 #[cfg(test)]
 mod test {
     use super::{PyFrozenSet, PySet};
-    use crate::{AsPyRef, ObjectProtocol, PyTryFrom, Python, ToPyObject};
+    use crate::{AsPyRef, IntoPy, ObjectProtocol, PyObject, PyTryFrom, Python, ToPyObject};
     use std::collections::{BTreeSet, HashSet};
     use std::iter::FromIterator;
 
@@ -451,4 +482,19 @@ mod test {
             BTreeSet::from_iter([1, 2, 3, 4, 5].iter().copied())
         );
     }
+
+    #[test]
+    fn test_set_into_py() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+
+        let bt: BTreeSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
+        let hs: HashSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
+
+        let bto: PyObject = bt.clone().into_py(py);
+        let hso: PyObject = hs.clone().into_py(py);
+
+        assert_eq!(bt, bto.extract(py).unwrap());
+        assert_eq!(hs, hso.extract(py).unwrap());
+    }
 }