1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

// Copyright (c) The nextest Contributors
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Support for string-only writes.
//!
//! There are potential situations where nextest needs to abstract over multiple kinds of writers,
//! but some of them do not accept arbitrary bytes -- they must be valid UTF-8.
//!
//! This is similar to [`std::fmt::Write`], but it returns [`std::io::Error`] instead for better
//! error handling.

use std::{fmt, io};

/// A trait that abstracts over writing strings to a writer.
///
/// For more, see the [module-level documentation](self).
pub trait WriteStr {
    /// Writes a string to the writer.
    fn write_str(&mut self, s: &str) -> io::Result<()>;

    /// Flushes the writer, ensuring that all intermediately buffered contents reach their
    /// destination.
    fn write_str_flush(&mut self) -> io::Result<()>;

    /// Writes a single character to the writer.
    fn write_char(&mut self, c: char) -> io::Result<()> {
        self.write_str(c.encode_utf8(&mut [0; 4]))
    }

    /// Writes a formatted string to the writer.
    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
        // This code is adapted from the `write_fmt` implementation for `std::io::Write`, and is
        // used under the terms of the MIT and Apache-2.0 licenses.

        // Create a shim which translates self to a fmt::Write and saves off errors instead of
        // discarding them.
        struct Adapter<'a, T: ?Sized> {
            inner: &'a mut T,
            error: Result<(), io::Error>,
        }

        impl<T: ?Sized + WriteStr> fmt::Write for Adapter<'_, T> {
            fn write_str(&mut self, s: &str) -> fmt::Result {
                match self.inner.write_str(s) {
                    Ok(()) => Ok(()),
                    Err(e) => {
                        self.error = Err(e);
                        Err(fmt::Error)
                    }
                }
            }
        }

        let mut output = Adapter {
            inner: self,
            error: Ok(()),
        };
        match fmt::write(&mut output, fmt) {
            Ok(()) => Ok(()),
            Err(_) => {
                // check if the error came from the underlying `Write` or not
                if output.error.is_err() {
                    output.error
                } else {
                    Err(io::Error::new(io::ErrorKind::Other, "formatter error"))
                }
            }
        }
    }
}

impl WriteStr for String {
    fn write_str(&mut self, s: &str) -> io::Result<()> {
        self.push_str(s);
        Ok(())
    }

    fn write_str_flush(&mut self) -> io::Result<()> {
        Ok(())
    }

    fn write_char(&mut self, c: char) -> io::Result<()> {
        self.push(c);
        Ok(())
    }
}

impl<T: WriteStr + ?Sized> WriteStr for &mut T {
    fn write_str(&mut self, s: &str) -> io::Result<()> {
        (**self).write_str(s)
    }

    fn write_str_flush(&mut self) -> io::Result<()> {
        (**self).write_str_flush()
    }

    fn write_char(&mut self, c: char) -> io::Result<()> {
        (**self).write_char(c)
    }

    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
        (**self).write_fmt(fmt)
    }
}

// Add more impls as needed.