pistonite_cu/process/pio/
read.rs

1use std::process::Stdio;
2
3use tokio::io::{AsyncBufRead, AsyncBufReadExt as _, AsyncReadExt as _};
4use tokio::process::{Child as TokioChild, ChildStderr, ChildStdout, Command as TokioCommand};
5
6use crate::{BoxedFuture, Context as _, co};
7
8use super::{ChildOutConfig, ChildOutTask};
9
10/// Buffer the output as a `Vec<u8>`
11///
12/// # Example
13/// ```rust
14/// # use pistonite_cu as cu;
15/// use cu::pre::*;
16///
17/// # #[cfg(unix)]
18/// # fn main() -> cu::Result<()> {
19/// let (hello, out) = cu::which("echo")?.command()
20///     .arg("Hello, world!")
21///     .stdout(cu::pio::buffer())
22///     .stdie_null()
23///     .spawn()?;
24///
25/// assert_eq!(b"Hello, world!\n".to_vec(), out.join()??);
26/// # Ok(()) }
27/// # #[cfg(not(unix))] fn main() {}
28/// ```
29#[inline(always)]
30pub fn buffer() -> Buffer {
31    Buffer
32}
33pub struct Buffer;
34impl ChildOutConfig for Buffer {
35    type Task = BufferTask;
36    type __Null = super::__OCNonNull;
37    fn configure_stdout(&mut self, command: &mut TokioCommand) {
38        command.stdout(Stdio::piped());
39    }
40    fn configure_stderr(&mut self, command: &mut TokioCommand) {
41        command.stderr(Stdio::piped());
42    }
43    fn take(
44        self,
45        child: &mut TokioChild,
46        _: Option<&str>,
47        is_out: bool,
48    ) -> crate::Result<Self::Task> {
49        Ok(BufferTask(super::take_child_out(child, is_out)?))
50    }
51}
52
53pub struct BufferTask(Result<ChildStdout, ChildStderr>);
54impl ChildOutTask for BufferTask {
55    type Output = co::Handle<crate::Result<Vec<u8>>>;
56
57    fn run(self) -> (Option<BoxedFuture<()>>, Self::Output) {
58        // we use a channel instead of returning the future in the output
59        // slot because we want to drive it internally to be less error-prone.
60        // if we return the future in the second slot without `spawn`,
61        // there could be a scenario where the future is never polled,
62        // resulting in stdout being stuck
63        let (send, recv) = tokio::sync::oneshot::channel();
64        let output: Self::Output =
65            co::spawn(async move { recv.await.context("failed to receive output as bytes")? });
66        let task = async move {
67            let buf = read_to_end(self.0).await;
68            let _: Result<_, _> = send.send(buf);
69        };
70        (Some(Box::pin(task)), output)
71    }
72}
73
74/// Buffer the output as a `String`. Will error if the output is not valid UTF-8.
75///
76/// # Example
77/// ```rust
78/// # use pistonite_cu as cu;
79/// use cu::pre::*;
80///
81/// # #[cfg(unix)]
82/// # fn main() -> cu::Result<()> {
83/// let (hello, out) = cu::which("echo")?.command()
84///     .arg("Hello, world!")
85///     .stdout(cu::pio::string())
86///     .stdie_null()
87///     .spawn()?;
88///
89/// assert_eq!("Hello, world!\n", out.join()??);
90/// # Ok(()) }
91/// # #[cfg(not(unix))] fn main() {}
92/// ```
93#[inline(always)]
94pub fn string() -> BufferString {
95    BufferString
96}
97pub struct BufferString;
98impl ChildOutConfig for BufferString {
99    type Task = BufferStringTask;
100    type __Null = super::__OCNonNull;
101    fn configure_stdout(&mut self, command: &mut TokioCommand) {
102        command.stdout(Stdio::piped());
103    }
104    fn configure_stderr(&mut self, command: &mut TokioCommand) {
105        command.stderr(Stdio::piped());
106    }
107    fn take(
108        self,
109        child: &mut TokioChild,
110        _: Option<&str>,
111        is_out: bool,
112    ) -> crate::Result<Self::Task> {
113        Ok(BufferStringTask(super::take_child_out(child, is_out)?))
114    }
115}
116
117pub struct BufferStringTask(Result<ChildStdout, ChildStderr>);
118impl ChildOutTask for BufferStringTask {
119    type Output = co::Handle<crate::Result<String>>;
120
121    fn run(self) -> (Option<BoxedFuture<()>>, Self::Output) {
122        // we use a channel instead of returning the future in the output
123        // slot because we want to drive it internally to be less error-prone.
124        // if we return the future in the second slot without `spawn`,
125        // there could be a scenario where the future is never polled,
126        // resulting in stdout being stuck
127        let (send, recv) = tokio::sync::oneshot::channel();
128        let output: Self::Output =
129            co::spawn(async move { recv.await.context("failed to receive output as bytes")? });
130        let task = async move {
131            let string = read_to_end(self.0).await.and_then(|buf| {
132                String::from_utf8(buf).context("failed to decode child output as utf-8")
133            });
134            let _: Result<_, _> = send.send(string);
135        };
136        (Some(Box::pin(task)), output)
137    }
138}
139
140async fn read_to_end(r: Result<ChildStdout, ChildStderr>) -> crate::Result<Vec<u8>> {
141    let mut buf = Vec::new();
142    match r {
143        Ok(mut r) => r.read_to_end(&mut buf).await,
144        Err(mut r) => r.read_to_end(&mut buf).await,
145    }
146    .context("io error while reading output")?;
147    Ok(buf)
148}
149
150/// Create reader to access output line-by-line.
151///
152/// The items will be available for read as soon as they are piped from the child,
153/// without having to wait for the child to finish first.
154///
155/// Note that for small outputs, it's usually better to simply use [`string()`]
156/// and call `.lines()` on the result.
157///
158/// # Example
159/// ```rust
160/// # use pistonite_cu as cu;
161/// use cu::pre::*;
162///
163/// # #[cfg(unix)]
164/// # fn main() -> cu::Result<()> {
165/// let (child, mut lines) = cu::which("bash")?.command()
166///     .args(["-c", r#"for i in {1..5}; do echo "Line $i"; done"# ])
167///     .stdout(cu::pio::lines())
168///     .stdie_null()
169///     .spawn()?;
170///
171/// assert_eq!(lines.next().unwrap()?, "Line 1");
172/// assert_eq!(lines.next().unwrap()?, "Line 2");
173/// assert_eq!(lines.next().unwrap()?, "Line 3");
174/// assert_eq!(lines.next().unwrap()?, "Line 4");
175/// assert_eq!(lines.next().unwrap()?, "Line 5");
176/// assert!(lines.next().is_none());
177///
178/// child.wait_nz()?;
179/// # Ok(())}
180/// # #[cfg(not(unix))] fn main() {}
181/// ```
182///
183/// # Blocking
184/// Since the iterator blocks the thread while waiting for the next line
185/// to be available, use [`co_lines`] when in an async context.
186#[inline(always)]
187pub fn lines() -> Lines {
188    Lines
189}
190pub struct Lines;
191impl ChildOutConfig for Lines {
192    type Task = LinesTask;
193    type __Null = super::__OCNonNull;
194    fn configure_stdout(&mut self, command: &mut TokioCommand) {
195        command.stdout(Stdio::piped());
196    }
197    fn configure_stderr(&mut self, command: &mut TokioCommand) {
198        command.stderr(Stdio::piped());
199    }
200    fn take(
201        self,
202        child: &mut TokioChild,
203        _: Option<&str>,
204        is_out: bool,
205    ) -> crate::Result<Self::Task> {
206        Ok(LinesTask(super::take_child_out(child, is_out)?))
207    }
208}
209
210pub struct LinesTask(Result<ChildStdout, ChildStderr>);
211impl ChildOutTask for LinesTask {
212    type Output = LinesOutput;
213
214    fn run(self) -> (Option<BoxedFuture<()>>, Self::Output) {
215        let (send, recv) = std::sync::mpsc::channel();
216        let output = LinesOutput { recv, done: false };
217        (Some(Box::pin(self.main(send))), output)
218    }
219}
220
221impl LinesTask {
222    async fn main(self, send: std::sync::mpsc::Sender<Option<crate::Result<String>>>) {
223        match self.0 {
224            Ok(r) => read_send_line(tokio::io::BufReader::new(r), send).await,
225            Err(r) => read_send_line(tokio::io::BufReader::new(r), send).await,
226        }
227    }
228}
229
230/// Output of [`lines`], a synchronous iterator that can reads the output line-by-line.
231///
232/// Just like the standard library's `lines()` iterator, the items have no line endings.
233pub struct LinesOutput {
234    recv: std::sync::mpsc::Receiver<Option<crate::Result<String>>>,
235    done: bool,
236}
237impl Iterator for LinesOutput {
238    type Item = crate::Result<String>;
239
240    fn next(&mut self) -> Option<Self::Item> {
241        if self.done {
242            return None;
243        }
244        match self.recv.recv() {
245            Ok(Some(x)) => {
246                if x.is_err() {
247                    self.done = true;
248                }
249                Some(x)
250            }
251            // no more data
252            _ => {
253                self.done = true;
254                None
255            }
256        }
257    }
258}
259
260/// Create reader to access output line-by-line.
261///
262/// The items will be available for read as soon as they are piped from the child,
263/// without having to wait for the child to finish first.
264///
265/// Note that for small outputs, it's usually better to simply use [`string()`]
266/// and call `.lines()` on the result.
267///
268/// # Example
269/// ```rust
270/// # use pistonite_cu as cu;
271/// use cu::pre::*;
272///
273/// # #[cfg(unix)]
274/// # #[cu::cli]
275/// # async fn main(_: cu::cli::Flags) -> cu::Result<()> {
276/// let (child, mut lines) = cu::which("bash")?.command()
277///     .args(["-c", r#"for i in {1..5}; do echo "Line $i"; done"# ])
278///     .stdout(cu::pio::co_lines())
279///     .stdie_null()
280///     .co_spawn().await?;
281///
282/// let mut i = 1;
283/// while let Some(line) = lines.next().await {
284///     let line = line?;
285///     assert_eq!(line, format!("Line {i}"));
286///     i+=1;
287/// }
288///
289/// child.co_wait_nz().await?;
290/// # Ok(())}
291/// # #[cfg(not(unix))] fn main() {}
292/// ```
293#[inline(always)]
294pub fn co_lines() -> CoLines {
295    CoLines
296}
297pub struct CoLines;
298impl ChildOutConfig for CoLines {
299    type Task = CoLinesTask;
300    type __Null = super::__OCNonNull;
301    fn configure_stdout(&mut self, command: &mut TokioCommand) {
302        command.stdout(Stdio::piped());
303    }
304    fn configure_stderr(&mut self, command: &mut TokioCommand) {
305        command.stderr(Stdio::piped());
306    }
307    fn take(
308        self,
309        child: &mut TokioChild,
310        _: Option<&str>,
311        is_out: bool,
312    ) -> crate::Result<Self::Task> {
313        Ok(CoLinesTask(super::take_child_out(child, is_out)?))
314    }
315}
316pub struct CoLinesTask(Result<ChildStdout, ChildStderr>);
317impl ChildOutTask for CoLinesTask {
318    type Output = CoLinesOutput;
319
320    fn run(self) -> (Option<BoxedFuture<()>>, Self::Output) {
321        let (send, recv) = tokio::sync::mpsc::unbounded_channel();
322        let output = CoLinesOutput { recv, done: false };
323        (Some(Box::pin(self.main(send))), output)
324    }
325}
326
327impl CoLinesTask {
328    async fn main(self, send: tokio::sync::mpsc::UnboundedSender<Option<crate::Result<String>>>) {
329        match self.0 {
330            Ok(r) => read_send_line(tokio::io::BufReader::new(r), send).await,
331            Err(r) => read_send_line(tokio::io::BufReader::new(r), send).await,
332        }
333    }
334}
335
336/// Output of [`co_lines`], an ansynchronous iterator that can reads the output line-by-line.
337///
338/// Just like the standard library's `lines()` iterator, the items have no line endings.
339pub struct CoLinesOutput {
340    recv: tokio::sync::mpsc::UnboundedReceiver<Option<crate::Result<String>>>,
341    done: bool,
342}
343impl CoLinesOutput {
344    /// Get the next line
345    pub async fn next(&mut self) -> Option<crate::Result<String>> {
346        if self.done {
347            return None;
348        }
349        match self.recv.recv().await {
350            Some(Some(x)) => {
351                if x.is_err() {
352                    self.done = true;
353                }
354                Some(x)
355            }
356            // no more data
357            _ => {
358                self.done = true;
359                None
360            }
361        }
362    }
363}
364
365trait LineSender {
366    fn send(&self, payload: Option<crate::Result<String>>) -> bool;
367}
368impl LineSender for std::sync::mpsc::Sender<Option<crate::Result<String>>> {
369    fn send(&self, payload: Option<crate::Result<String>>) -> bool {
370        self.send(payload).is_ok()
371    }
372}
373impl LineSender for tokio::sync::mpsc::UnboundedSender<Option<crate::Result<String>>> {
374    fn send(&self, payload: Option<crate::Result<String>>) -> bool {
375        self.send(payload).is_ok()
376    }
377}
378async fn read_send_line<R: AsyncBufRead + Unpin, S: LineSender>(r: R, send: S) {
379    let mut lines = r.lines();
380    loop {
381        match lines.next_line().await {
382            Err(e) => {
383                // io error
384                let payload = Err(e).context("error reading line from child output");
385                send.send(Some(payload));
386                return;
387            }
388            Ok(Some(x)) => {
389                if !send.send(Some(Ok(x))) {
390                    // try sending the end message
391                    send.send(None);
392                    return;
393                }
394            }
395            Ok(None) => {
396                // no more lines
397                let _ = send.send(None);
398                return;
399            }
400        }
401    }
402}
pistonite_cu/process/pio/read.rs

pistonite_cu/process/pio/
read.rs
