[Question]: How to use `asyncio`'s stream types as `tokio`'s AsyncRead/AsyncWrite?

[JadKHaddad]

• How to use StreamReader as a generic AsyncRead?
• How to use StreamWriter as a generic AsyncWrite?

I took a look at the unstable feature unstable-streams to convert a stream of bytes into an AsyncRead, which looks roughly like this:

fn internal_function_takes_async_read<T: AsyncRead>(reader: T) {}


#[pyclass]
pub struct AsyncGenBytes {
    /// An async generator that produces byte arrays
    py_async_gen: Py<PyAny>,
}

// Does not really compile because AsyncGenBytes does not implement PyFunctionArgument (not important for this example)
#[pyfunction]
fn takes_async_gen_bytes_exposed_to_python(py_async_gen: AsyncGenBytes) -> PyResult<()> {
    let stream = Python::with_gil(|py| {
        let bound = py_async_gen.py_async_gen.into_bound(py);

        pyo3_async_runtimes::tokio::into_stream_v1(bound)
    })?;

    let byte_stream = stream.then(|res| async {
        Python::with_gil(|py| match res {
            Ok(obj) => match obj.extract::<Vec<u8>>(py) {
                Ok(vec) => Ok(Bytes::from(vec)),
                Err(e) => Err(std::io::Error::other(format!(
                    "Failed to extract bytes: {e}"
                ))),
            },
            Err(e) => Err(std::io::Error::other(format!("Python error: {e}"))),
        })
    });
    
    // StreamReader implements AsyncRead
    let stream_reader = tokio_util::io::StreamReader::new(byte_stream);

    // do the things you want to do
    internal_function_takes_async_read(stream_reader);

    Ok(())
}

This seems to me like hacking around to get things done.

Suggestions on how to do it the right way are very appreciated.

Thanks all for this GREAT project!

[JadKHaddad]

I came up with this

async fn internal_function_takes_async_read<T: AsyncRead>(reader: T) {}

#[pyclass]
#[derive(Clone)]
pub struct AsyncReader {
    /// any asyncio.StreamReader
    asyncio_stream_reader: Arc<PyObject>,
}

#[pymethods]
impl AsyncReader {
    #[new]
    fn new(asyncio_stream_reader: PyObject) -> Self {
        Self {
            asyncio_stream_reader: Arc::new(asyncio_stream_reader),
        }
    }
}

#[pyfunction]
#[pyo3(signature=(read))]
pub fn takes_async_read_exposed_to_python(
    py: Python<'_>,
    read: AsyncReader,
) -> PyResult<Bound<PyAny>> {
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
        let read = read.async_read();

        internal_function_takes_async_read(read).await;

        Ok(())
    })
}

impl AsyncReader {
    fn async_read(&self) -> impl AsyncRead {
        tokio_util::io::StreamReader::new(self.unfold())
    }

    fn unfold(&self) -> impl futures::Stream<Item = Result<Bytes, std::io::Error>> + Send {
        let reader = self.clone();

        futures::stream::unfold(reader, |reader| async move {
            let py_fut_result = Python::with_gil(|py| {
                let awaitable = reader
                    .asyncio_stream_reader
                    .call_method1(py, "read", (1024,))?;

                pyo3_async_runtimes::tokio::into_future(awaitable.into_bound(py))
            });

            let fut = match py_fut_result {
                Ok(fut) => fut,
                Err(err) => {
                    return Some((Err(std::io::Error::other(err.to_string())), reader));
                }
            };

            let result = fut.await;

            let extracted = Python::with_gil(|py| match result {
                Ok(obj) => match obj.extract::<Vec<u8>>(py) {
                    Ok(data) => Ok(data),
                    Err(e) => Err(std::io::Error::new(
                        std::io::ErrorKind::InvalidData,
                        e.to_string(),
                    )),
                },
                Err(e) => Err(std::io::Error::other(e.to_string())),
            });

            match extracted {
                Ok(bytes) if !bytes.is_empty() => Some((Ok(Bytes::from(bytes)), reader)),
                Ok(_) => {
                    // EOF
                    None
                }
                Err(e) => Some((Err(e), reader)),
            }
        })
    }
}

[JadKHaddad]

Latest version I am using to convert into AsyncRead + AsyncWrite. Might be helpful for somebody looking for something similar:

use std::future::Future;

use bytes::Bytes;
use futures::{Sink, Stream};
use pyo3::{types::PyBytes, PyErr, PyObject, PyResult, Python};
use tokio::io::{AsyncRead, AsyncWrite};

pub trait IO {
    fn into_tokio_async_read_and_write(self) -> impl AsyncRead + AsyncWrite;
}

impl IO for (PyObject, PyObject) {
    // Converts a tuple of Python asyncio StreamReader and StreamWriter to a Tokio AsyncRead and AsyncWrite.
    fn into_tokio_async_read_and_write(self) -> impl AsyncRead + AsyncWrite {
        tokio::io::join(
            async_io_stream_reader_to_async_read(self.0),
            async_io_stream_writer_to_async_write(self.1),
        )
    }
}

/// Converts a Python [asyncio.StreamReader](https://docs.python.org/3/library/asyncio-stream.html#streamreader) to a [`Stream`].
fn async_io_stream_reader_to_stream(reader: PyObject) -> impl Stream<Item = PyResult<Bytes>> {
    futures::stream::unfold(reader, |reader| async move {
        let extract = async {
            let fut = Python::with_gil(|py| {
                let awaitable = reader.call_method1(py, "read", (1024,))?;

                pyo3_async_runtimes::tokio::into_future(awaitable.into_bound(py))
            })?;

            let obj = fut.await?;

            let extracted = Python::with_gil(|py| obj.extract::<Vec<u8>>(py))?;

            PyResult::Ok(extracted)
        };

        match extract.await {
            Ok(bytes) if !bytes.is_empty() => Some((Ok(Bytes::from(bytes)), reader)),
            Ok(_) => {
                // EOF, terminate stream
                None
            }
            Err(err) => Some((Err(err), reader)),
        }
    })
}

/// Converts a Python [asyncio.StreamReader](https://docs.python.org/3/library/asyncio-stream.html#streamreader) to an [`AsyncRead`].
fn async_io_stream_reader_to_async_read(stream_reader: PyObject) -> impl AsyncRead {
    tokio_util::io::StreamReader::new(async_io_stream_reader_to_stream(stream_reader))
}

/// Converts a Python [asyncio.StreamWriter](https://docs.python.org/3/library/asyncio-stream.html#streamwriter) to a [`Sink`].
fn async_io_stream_writer_to_sink(writer: PyObject) -> impl for<'a> Sink<&'a [u8], Error = PyErr> {
    fn f(writer: PyObject, data: &[u8]) -> impl Future<Output = Result<PyObject, PyErr>> + 'static {
        let write_result =
            Python::with_gil(|py| writer.call_method1(py, "write", (PyBytes::new(py, data),)));

        async {
            write_result?;

            let fut = Python::with_gil(|py| {
                let awaitable = writer.call_method0(py, "drain")?;

                pyo3_async_runtimes::tokio::into_future(awaitable.into_bound(py))
            })?;

            fut.await?;

            Ok(writer)
        }
    }

    futures::sink::unfold(writer, f)
}

/// Converts a Python [asyncio.StreamWriter](https://docs.python.org/3/library/asyncio-stream.html#streamwriter) to an [`AsyncWrite`].
fn async_io_stream_writer_to_async_write(stream_writer: PyObject) -> impl AsyncWrite {
    tokio_util::io::SinkWriter::new(async_io_stream_writer_to_sink(stream_writer))
}

Now one could use these types inside a tokio spawned task

#[pyo3::pyfunction]
#[pyo3(signature=(read, write))]
fn function<'p>(py: Python<'p>, read: PyObject, write: PyObject) -> PyResult<Bound<'p, PyAny>> {
    let read_write = (read, write).into_tokio_async_read_and_write();

    future_into_py(py, async move {
        let task_locals = Python::with_gil(pyo3_async_runtimes::tokio::get_current_locals)?;
        tokio::spawn(pyo3_async_runtimes::tokio::scope(task_locals, async move {
            // Do async read and write stuff
            let mut read_write = Box::pin(read_write);
            
            let mut buf: [u8; 1024] = [0; 1024];

            read_write.read(&mut buf).await.unwrap();

            read_write.write_all(b"hi").await.unwrap();
        }));

        Ok(())
    })
}

Now in Python I use it like this

import lib.function
import asyncio

async def main():
    read, write = await asyncio.open_connection("127.0.0.1", 2775)
    
    await function(read, write)

if __name__ == "__main__":
    asyncio.run(main())