partial_io/
write.rs

1// Copyright (c) The partial-io Contributors
2// SPDX-License-Identifier: MIT
3
4//! This module contains a writer wrapper that breaks writes up according to a
5//! provided iterator.
6
7use std::{
8    cmp, fmt,
9    io::{self, Read, Write},
10};
11
12use crate::{make_ops, PartialOp};
13
14/// A writer wrapper that breaks inner `Write` instances up according to the
15/// provided iterator.
16///
17/// # Examples
18///
19/// ```rust
20/// use std::io::Write;
21///
22/// use partial_io::{PartialOp, PartialWrite};
23///
24/// let writer = Vec::new();
25/// let iter = ::std::iter::repeat(PartialOp::Limited(1));
26/// let mut partial_writer = PartialWrite::new(writer, iter);
27/// let in_data = vec![1, 2, 3, 4];
28///
29/// let size = partial_writer.write(&in_data).unwrap();
30/// assert_eq!(size, 1);
31/// assert_eq!(&partial_writer.get_ref()[..], &[1]);
32/// ```
33pub struct PartialWrite<W> {
34    inner: W,
35    ops: Box<dyn Iterator<Item = PartialOp> + Send>,
36}
37
38impl<W> PartialWrite<W>
39where
40    W: Write,
41{
42    /// Creates a new `PartialWrite` wrapper over the writer with the specified `PartialOp`s.
43    pub fn new<I>(inner: W, iter: I) -> Self
44    where
45        I: IntoIterator<Item = PartialOp> + 'static,
46        I::IntoIter: Send,
47    {
48        PartialWrite {
49            inner,
50            // Use fuse here so that we don't keep calling the inner iterator
51            // once it's returned None.
52            ops: make_ops(iter),
53        }
54    }
55
56    /// Sets the `PartialOp`s for this writer.
57    pub fn set_ops<I>(&mut self, iter: I) -> &mut Self
58    where
59        I: IntoIterator<Item = PartialOp> + 'static,
60        I::IntoIter: Send,
61    {
62        self.ops = make_ops(iter);
63        self
64    }
65
66    /// Acquires a reference to the underlying writer.
67    pub fn get_ref(&self) -> &W {
68        &self.inner
69    }
70
71    /// Acquires a mutable reference to the underlying writer.
72    pub fn get_mut(&mut self) -> &mut W {
73        &mut self.inner
74    }
75
76    /// Consumes this wrapper, returning the underlying writer.
77    pub fn into_inner(self) -> W {
78        self.inner
79    }
80}
81
82impl<W> Write for PartialWrite<W>
83where
84    W: Write,
85{
86    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
87        match self.ops.next() {
88            Some(PartialOp::Limited(n)) => {
89                let len = cmp::min(n, buf.len());
90                self.inner.write(&buf[..len])
91            }
92            Some(PartialOp::Err(err)) => Err(io::Error::new(
93                err,
94                "error during write, generated by partial-io",
95            )),
96            Some(PartialOp::Unlimited) | None => self.inner.write(buf),
97        }
98    }
99
100    fn flush(&mut self) -> io::Result<()> {
101        match self.ops.next() {
102            Some(PartialOp::Err(err)) => Err(io::Error::new(
103                err,
104                "error during flush, generated by partial-io",
105            )),
106            _ => self.inner.flush(),
107        }
108    }
109}
110
111// Forwarding impl to support duplex structs.
112impl<W> Read for PartialWrite<W>
113where
114    W: Read + Write,
115{
116    #[inline]
117    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
118        self.inner.read(buf)
119    }
120}
121
122impl<W> fmt::Debug for PartialWrite<W>
123where
124    W: fmt::Debug,
125{
126    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
127        f.debug_struct("PartialWrite")
128            .field("inner", &self.inner)
129            .finish()
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    use std::fs::File;
138
139    use crate::tests::assert_send;
140
141    #[test]
142    fn test_sendable() {
143        assert_send::<PartialWrite<File>>();
144    }
145}