1 use crate::cli::{IsTerminal, StdoutStream}; 2 use crate::p2; 3 use bytes::Bytes; 4 use std::io::{self, Write}; 5 use std::pin::Pin; 6 use std::task::{Context, Poll}; 7 use tokio::io::AsyncWrite; 8 use wasmtime_wasi_io::streams::OutputStream; 9 10 // Implementation for tokio::io::Stdout 11 impl IsTerminal for tokio::io::Stdout { 12 fn is_terminal(&self) -> bool { 13 std::io::stdout().is_terminal() 14 } 15 } 16 impl 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 26 impl IsTerminal for std::io::Stdout { 27 fn is_terminal(&self) -> bool { 28 std::io::IsTerminal::is_terminal(self) 29 } 30 } 31 impl 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 41 impl IsTerminal for tokio::io::Stderr { 42 fn is_terminal(&self) -> bool { 43 std::io::stderr().is_terminal() 44 } 45 } 46 impl 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 56 impl IsTerminal for std::io::Stderr { 57 fn is_terminal(&self) -> bool { 58 std::io::IsTerminal::is_terminal(self) 59 } 60 } 61 impl 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 70 enum StdioOutputStream { 71 Stdout, 72 Stderr, 73 } 74 75 impl 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 97 impl 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] 120 impl p2::Pollable for StdioOutputStream { 121 async fn ready(&mut self) {} 122 } 123