1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

#![cfg(not(Py_LIMITED_API))]

//! Objects related to PyBuffer and PyStr
use pyo3::buffer::PyBuffer;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyMemoryView, PyString};

/// This is for confirming that PyBuffer does not cause memory leak
#[pyclass]
struct BytesExtractor {}

#[pymethods]
impl BytesExtractor {
    #[new]
    pub fn __new__() -> Self {
        BytesExtractor {}
    }

    #[staticmethod]
    pub fn from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<usize> {
        let byte_vec: Vec<u8> = bytes.extract()?;
        Ok(byte_vec.len())
    }

    #[staticmethod]
    pub fn from_str(string: &Bound<'_, PyString>) -> PyResult<usize> {
        let rust_string: String = string.extract()?;
        Ok(rust_string.len())
    }

    #[staticmethod]
    pub fn from_str_lossy(string: &Bound<'_, PyString>) -> usize {
        let rust_string_lossy: String = string.to_string_lossy().to_string();
        rust_string_lossy.len()
    }

    #[staticmethod]
    pub fn from_buffer(buf: &Bound<'_, PyAny>) -> PyResult<usize> {
        let buf = PyBuffer::<u8>::get_bound(buf)?;
        Ok(buf.item_count())
    }
}

#[pyfunction]
fn return_memoryview(py: Python<'_>) -> PyResult<Bound<'_, PyMemoryView>> {
    let bytes = PyBytes::new_bound(py, b"hello world");
    PyMemoryView::from_bound(&bytes)
}

#[pymodule]
pub fn buf_and_str(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<BytesExtractor>()?;
    m.add_function(wrap_pyfunction!(return_memoryview, m)?)?;
    Ok(())
}