strftime/format/
write.rs

1//! This module is a copy of the [`std::io::Write`] implementation, in order to
2//! use it in a no-std context.
3//!
4//! [`std::io::Write`]: <https://doc.rust-lang.org/std/io/trait.Write.html>
5
6#[cfg(feature = "alloc")]
7use alloc::vec::Vec;
8use core::fmt;
9use core::str;
10
11use crate::Error;
12
13/// An `Adapter` implements [`core::fmt::Write`] from a [`Write`] object,
14/// storing write errors instead of discarding them.
15struct Adapter<'a, T: ?Sized> {
16    /// Inner writer.
17    inner: &'a mut T,
18    /// Write result.
19    error: Result<(), Error>,
20}
21
22impl<T: Write + ?Sized> fmt::Write for Adapter<'_, T> {
23    fn write_str(&mut self, s: &str) -> fmt::Result {
24        match self.inner.write_all(s.as_bytes()) {
25            Ok(()) => Ok(()),
26            Err(e) => {
27                self.error = Err(e);
28                Err(fmt::Error)
29            }
30        }
31    }
32}
33
34/// Simplified copy of the [`std::io::Write`] trait.
35///
36/// [`std::io::Write`]: <https://doc.rust-lang.org/std/io/trait.Write.html>
37pub(crate) trait Write {
38    /// Write a buffer into this writer, returning how many bytes were written.
39    fn write(&mut self, data: &[u8]) -> Result<usize, Error>;
40
41    /// Attempts to write an entire buffer into this writer.
42    fn write_all(&mut self, mut data: &[u8]) -> Result<(), Error> {
43        while !data.is_empty() {
44            match self.write(data)? {
45                0 => return Err(Error::WriteZero),
46                n => data = &data[n..],
47            }
48        }
49        Ok(())
50    }
51
52    /// Writes a formatted string into this writer, returning any error
53    /// encountered.
54    fn write_fmt(&mut self, fmt_args: fmt::Arguments<'_>) -> Result<(), Error> {
55        let mut output = Adapter {
56            inner: self,
57            error: Ok(()),
58        };
59
60        match fmt::write(&mut output, fmt_args) {
61            Ok(()) => Ok(()),
62            Err(_) if output.error.is_err() => output.error,
63            Err(err) => Err(err.into()),
64        }
65    }
66}
67
68/// Write is implemented for `&mut [u8]` by copying into the slice, overwriting
69/// its data.
70impl Write for &mut [u8] {
71    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
72        let size = data.len().min(self.len());
73        let (a, b) = core::mem::take(self).split_at_mut(size);
74        a.copy_from_slice(&data[..size]);
75        *self = b;
76        Ok(size)
77    }
78}
79
80/// Wrapper for a [`core::fmt::Write`] writer.
81pub(crate) struct FmtWrite<'a> {
82    /// Inner writer.
83    inner: &'a mut dyn fmt::Write,
84}
85
86impl<'a> FmtWrite<'a> {
87    /// Construct a new `FmtWrite`.
88    pub(crate) fn new(inner: &'a mut dyn fmt::Write) -> Self {
89        Self { inner }
90    }
91}
92
93/// Write is implemented for `FmtWrite` by writing to its inner writer.
94impl Write for FmtWrite<'_> {
95    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
96        let data = str::from_utf8(data).expect("FmtWrite should only receive UTF-8 data");
97        self.inner.write_str(data)?;
98        Ok(data.len())
99    }
100
101    fn write_fmt(&mut self, fmt_args: fmt::Arguments<'_>) -> Result<(), Error> {
102        Ok(self.inner.write_fmt(fmt_args)?)
103    }
104}
105
106/// Write is implemented for `Vec<u8>` by appending to the vector, growing as
107/// needed.
108#[cfg(feature = "alloc")]
109#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
110impl Write for Vec<u8> {
111    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
112        self.try_reserve(data.len())?;
113        self.extend_from_slice(data);
114        Ok(data.len())
115    }
116}
117
118/// Wrapper for a [`std::io::Write`] writer.
119#[cfg(feature = "std")]
120#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
121pub(crate) struct IoWrite<'a> {
122    /// Inner writer.
123    inner: &'a mut dyn std::io::Write,
124}
125
126#[cfg(feature = "std")]
127impl<'a> IoWrite<'a> {
128    /// Construct a new `IoWrite`.
129    pub(crate) fn new(inner: &'a mut dyn std::io::Write) -> Self {
130        Self { inner }
131    }
132}
133
134/// Write is implemented for `IoWrite` by writing to its inner writer.
135#[cfg(feature = "std")]
136impl Write for IoWrite<'_> {
137    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
138        Ok(self.inner.write(data)?)
139    }
140
141    fn write_all(&mut self, data: &[u8]) -> Result<(), Error> {
142        Ok(self.inner.write_all(data)?)
143    }
144
145    fn write_fmt(&mut self, fmt_args: fmt::Arguments<'_>) -> Result<(), Error> {
146        Ok(self.inner.write_fmt(fmt_args)?)
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn test_fmt_error() {
156        use core::fmt;
157
158        struct S;
159
160        impl fmt::Display for S {
161            fn fmt(&self, _: &mut fmt::Formatter<'_>) -> fmt::Result {
162                Err(fmt::Error)
163            }
164        }
165
166        let result = write!(&mut &mut [0u8; 1][..], "{S}");
167        assert!(matches!(result, Err(Error::FmtError(_))));
168    }
169
170    #[cfg(feature = "std")]
171    #[test]
172    fn test_io_write() {
173        let mut buf = Vec::new();
174
175        let mut writer = IoWrite::new(&mut buf);
176        writer.write_all(b"ok").unwrap();
177        write!(writer, "{}", 1).unwrap();
178
179        assert_eq!(buf, *b"ok1");
180    }
181
182    #[cfg(feature = "alloc")]
183    #[test]
184    fn test_fmt_write() {
185        use alloc::string::String;
186
187        let mut buf = String::new();
188        write!(FmtWrite::new(&mut buf), "{}", 1).unwrap();
189        assert_eq!(buf, "1");
190    }
191}