foundry_common/io/
stdin.rs

1//! Utility functions for reading from [`stdin`](std::io::stdin).
2
3use eyre::Result;
4use std::{
5    error::Error as StdError,
6    io::{self, BufRead, Read},
7    str::FromStr,
8};
9
10/// Unwraps the given `Option<T>` or [reads stdin into a String](read) and parses it as `T`.
11pub fn unwrap<T>(value: Option<T>, read_line: bool) -> Result<T>
12where
13    T: FromStr,
14    T::Err: StdError + Send + Sync + 'static,
15{
16    match value {
17        Some(value) => Ok(value),
18        None => parse(read_line),
19    }
20}
21
22/// Shortcut for `(unwrap(a), unwrap(b))`.
23pub fn unwrap2<A, B>(a: Option<A>, b: Option<B>) -> Result<(A, B)>
24where
25    A: FromStr,
26    B: FromStr,
27    A::Err: StdError + Send + Sync + 'static,
28    B::Err: StdError + Send + Sync + 'static,
29{
30    match (a, b) {
31        (Some(a), Some(b)) => Ok((a, b)),
32        (a, b) => Ok((unwrap(a, true)?, unwrap(b, true)?)),
33    }
34}
35
36/// [Reads stdin into a String](read) and parses it as `Vec<T>` using whitespaces as delimiters if
37/// the given `Vec<T>` is empty.
38pub fn unwrap_vec<T>(mut value: Vec<T>) -> Result<Vec<T>>
39where
40    T: FromStr,
41    T::Err: StdError + Send + Sync + 'static,
42{
43    if value.is_empty() {
44        let s = read(false)?;
45        value = s.split_whitespace().map(FromStr::from_str).collect::<Result<Vec<T>, _>>()?;
46    }
47
48    Ok(value)
49}
50
51/// Short-hand for `unwrap(value, true)`.
52pub fn unwrap_line<T>(value: Option<T>) -> Result<T>
53where
54    T: FromStr,
55    T::Err: StdError + Send + Sync + 'static,
56{
57    unwrap(value, true)
58}
59
60/// Reads bytes from [`stdin`][io::stdin] into a String.
61///
62/// If `read_line` is true, stop at the first newline (the `0xA` byte).
63pub fn parse<T>(read_line: bool) -> Result<T>
64where
65    T: FromStr,
66    T::Err: StdError + Send + Sync + 'static,
67{
68    read(read_line).and_then(|s| s.parse().map_err(Into::into))
69}
70
71/// Short-hand for `parse(true)`.
72pub fn parse_line<T>() -> Result<T>
73where
74    T: FromStr,
75    T::Err: StdError + Send + Sync + 'static,
76{
77    parse(true)
78}
79
80/// Reads bytes from [`stdin`][io::stdin] into a String.
81///
82/// If `read_line` is true, stop at the first newline (the `0xA` byte).
83pub fn read(read_line: bool) -> Result<String> {
84    let bytes = read_bytes(read_line)?;
85
86    if read_line {
87        // SAFETY: [BufRead::read_line] appends into a String
88        Ok(unsafe { String::from_utf8_unchecked(bytes) })
89    } else {
90        String::from_utf8(bytes).map_err(Into::into)
91    }
92}
93
94/// Reads bytes from [`stdin`][io::stdin].
95///
96/// If `read_line` is true, read up to the first newline excluded (the `0xA` byte).
97pub fn read_bytes(read_line: bool) -> Result<Vec<u8>> {
98    let mut stdin = io::stdin().lock();
99
100    if read_line {
101        let mut buf = String::new();
102        stdin.read_line(&mut buf)?;
103        // remove the trailing newline
104        if let Some(b'\n') = buf.as_bytes().last() {
105            buf.pop();
106            if let Some(b'\r') = buf.as_bytes().last() {
107                buf.pop();
108            }
109        }
110        Ok(buf.into_bytes())
111    } else {
112        let mut buf = Vec::new();
113        stdin.read_to_end(&mut buf)?;
114        Ok(buf)
115    }
116}
foundry_common/io/stdin.rs

foundry_common/io/
stdin.rs
