Add a test to detect memory leak of bytes->vec conversion

examples/rustapi_module/requirements-dev.txt

@@ -2,3 +2,4 @@ pip>=19.1
 hypothesis>=3.55
 pytest>=3.5.0
 setuptools-rust>=0.10.2
+psutil>=5.6

examples/rustapi_module/setup.py

@@ -97,6 +97,7 @@ setup(
         make_rust_extension("rustapi_module.datetime"),
         make_rust_extension("rustapi_module.subclassing"),
         make_rust_extension("rustapi_module.test_dict"),
+        make_rust_extension("rustapi_module.buf_and_str"),
     ],
     install_requires=install_requires,
     tests_require=tests_require,

examples/rustapi_module/src/buf_and_str.rs

@@ -0,0 +1,31 @@
+//! Objects related to PyBuffer and PyStr
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyString};
+
+/// This is for confirming that PyBuffer does not cause memory leak
+#[pyclass]
+struct BytesExtractor {}
+
+#[pymethods]
+impl BytesExtractor {
+    #[new]
+    pub fn __new__(obj: &PyRawObject) {
+        obj.init({ BytesExtractor {} });
+    }
+
+    pub fn to_vec(&mut self, bytes: &PyBytes) -> PyResult<usize> {
+        let byte_vec: Vec<u8> = bytes.extract().unwrap();
+        Ok(byte_vec.len())
+    }
+
+    pub fn to_str(&mut self, bytes: &PyString) -> PyResult<usize> {
+        let byte_str: String = bytes.extract().unwrap();
+        Ok(byte_str.len())
+    }
+}
+
+#[pymodule]
+pub fn buf_and_str(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<BytesExtractor>()?;
+    Ok(())
+}

examples/rustapi_module/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod buf_and_str;
 pub mod datetime;
 pub mod dict_iter;
 pub mod othermod;

examples/rustapi_module/tests/test_buf_and_str.py

@@ -0,0 +1,30 @@
+import os
+import psutil
+from rustapi_module.buf_and_str import BytesExtractor
+
+
+def test_pybuffer_doesnot_leak_memory():
+    N = int(1e5)
+    extractor = BytesExtractor()
+    process = psutil.Process(os.getpid())
+
+    def memory_diff(f):
+        before = process.memory_info().rss
+        f()
+        after = process.memory_info().rss
+        return after - before
+
+    message_b = b'\\(-"-;) Praying that memory leak would not happen..'
+    message_s = '\\(-"-;) Praying that memory leak would not happen..'
+
+    def to_vec():
+        for i in range(N):
+            extractor.to_vec(message_b)
+
+    def to_str():
+        for i in range(N):
+            extractor.to_str(message_s)
+
+    mv = memory_diff(to_vec)
+    ms = memory_diff(to_str)
+    assert abs(mv - ms) < 1000

tests/test_buffer_protocol.rs

@@ -2,7 +2,7 @@ use pyo3::class::PyBufferProtocol;
 use pyo3::exceptions::BufferError;
 use pyo3::ffi;
 use pyo3::prelude::*;
-use pyo3::types::IntoPyDict;
+use pyo3::types::{IntoPyDict, PyAny};
 use std::ffi::CStr;
 use std::os::raw::{c_int, c_void};
 use std::ptr;
@@ -76,3 +76,28 @@ fn test_buffer() {
     let d = [("ob", t)].into_py_dict(py);
     py.run("assert bytes(ob) == b' 23'", None, Some(d)).unwrap();
 }
+
+fn byte_to_vec(py_byte_string: &PyAny) -> PyResult<Vec<u8>> {
+    py_byte_string.extract()
+}
+
+fn byte_to_str(py_byte_string: &PyAny) -> PyResult<String> {
+    py_byte_string.extract()
+}
+
+#[test]
+fn test_byte_conversion() {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+    let bytes = py.eval("b'I AM BYTE-STRING'", None, None).unwrap();
+    let answer = "I AM BYTE-STRING";
+    assert_eq!(byte_to_vec(bytes).unwrap(), answer.as_bytes().to_vec());
+    assert_eq!(
+        byte_to_str(bytes)
+            .map_err(|err| {
+                err.print_and_set_sys_last_vars(py);
+            })
+            .unwrap(),
+        answer.to_string()
+    );
+}