wasmtime_wasi/cli/
stdout.rs

1use crate::cli::{IsTerminal, StdoutStream};
2use crate::p2;
3use bytes::Bytes;
4use std::io::{self, Write};
5use std::pin::Pin;
6use std::task::{Context, Poll};
7use tokio::io::AsyncWrite;
8use wasmtime_wasi_io::streams::OutputStream;
9
10// Implementation for tokio::io::Stdout
11impl IsTerminal for tokio::io::Stdout {
12    fn is_terminal(&self) -> bool {
13        std::io::stdout().is_terminal()
14    }
15}
16impl StdoutStream for tokio::io::Stdout {
17    fn p2_stream(&self) -> Box<dyn OutputStream> {
18        Box::new(StdioOutputStream::Stdout)
19    }
20    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
21        Box::new(StdioOutputStream::Stdout)
22    }
23}
24
25// Implementation for std::io::Stdout
26impl IsTerminal for std::io::Stdout {
27    fn is_terminal(&self) -> bool {
28        std::io::IsTerminal::is_terminal(self)
29    }
30}
31impl StdoutStream for std::io::Stdout {
32    fn p2_stream(&self) -> Box<dyn OutputStream> {
33        Box::new(StdioOutputStream::Stdout)
34    }
35    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
36        Box::new(StdioOutputStream::Stdout)
37    }
38}
39
40// Implementation for tokio::io::Stderr
41impl IsTerminal for tokio::io::Stderr {
42    fn is_terminal(&self) -> bool {
43        std::io::stderr().is_terminal()
44    }
45}
46impl StdoutStream for tokio::io::Stderr {
47    fn p2_stream(&self) -> Box<dyn OutputStream> {
48        Box::new(StdioOutputStream::Stderr)
49    }
50    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
51        Box::new(StdioOutputStream::Stderr)
52    }
53}
54
55// Implementation for std::io::Stderr
56impl IsTerminal for std::io::Stderr {
57    fn is_terminal(&self) -> bool {
58        std::io::IsTerminal::is_terminal(self)
59    }
60}
61impl StdoutStream for std::io::Stderr {
62    fn p2_stream(&self) -> Box<dyn OutputStream> {
63        Box::new(StdioOutputStream::Stderr)
64    }
65    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
66        Box::new(StdioOutputStream::Stderr)
67    }
68}
69
70enum StdioOutputStream {
71    Stdout,
72    Stderr,
73}
74
75impl OutputStream for StdioOutputStream {
76    fn write(&mut self, bytes: Bytes) -> p2::StreamResult<()> {
77        match self {
78            StdioOutputStream::Stdout => std::io::stdout().write_all(&bytes),
79            StdioOutputStream::Stderr => std::io::stderr().write_all(&bytes),
80        }
81        .map_err(|e| p2::StreamError::LastOperationFailed(anyhow::anyhow!(e)))
82    }
83
84    fn flush(&mut self) -> p2::StreamResult<()> {
85        match self {
86            StdioOutputStream::Stdout => std::io::stdout().flush(),
87            StdioOutputStream::Stderr => std::io::stderr().flush(),
88        }
89        .map_err(|e| p2::StreamError::LastOperationFailed(anyhow::anyhow!(e)))
90    }
91
92    fn check_write(&mut self) -> p2::StreamResult<usize> {
93        Ok(1024 * 1024)
94    }
95}
96
97impl AsyncWrite for StdioOutputStream {
98    fn poll_write(
99        self: Pin<&mut Self>,
100        _cx: &mut Context<'_>,
101        buf: &[u8],
102    ) -> Poll<io::Result<usize>> {
103        Poll::Ready(match *self {
104            StdioOutputStream::Stdout => std::io::stdout().write(buf),
105            StdioOutputStream::Stderr => std::io::stderr().write(buf),
106        })
107    }
108    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
109        Poll::Ready(match *self {
110            StdioOutputStream::Stdout => std::io::stdout().flush(),
111            StdioOutputStream::Stderr => std::io::stderr().flush(),
112        })
113    }
114    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
115        Poll::Ready(Ok(()))
116    }
117}
118
119#[async_trait::async_trait]
120impl p2::Pollable for StdioOutputStream {
121    async fn ready(&mut self) {}
122}