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