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