implement Write trait for BytesMut per issue #425

src/lib.rs

@@ -115,3 +115,67 @@ fn abort() -> ! {
         panic!("abort");
     }
 }
+
+/// extension traits
+#[cfg(feature = "std")]
+pub mod ext {
+    /// extension trait for BytesMut
+    pub trait BytesMutExt {
+        /// extension trait for BytesMut
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize>;
+        /// extension trait for BytesMut
+        fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize>;
+        /// extension trait for BytesMut
+        fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()>;
+        /// extension trait for BytesMut
+        fn flush(&mut self) -> std::io::Result<()>;
+    }
+    impl BytesMutExt for super::BytesMut {
+        #[inline]
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            self.extend_from_slice(buf);
+            Ok(buf.len())
+        }
+
+        #[inline]
+        fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
+            let len = bufs.iter().map(|b| b.len()).sum();
+            self.reserve(len);
+            for buf in bufs {
+                self.extend_from_slice(buf);
+            }
+            Ok(len)
+        }
+        #[inline]
+        fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
+            self.extend_from_slice(buf);
+            Ok(())
+        }
+
+        #[inline]
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    impl std::io::Write for dyn BytesMutExt {
+        #[inline]
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            self.write(buf)
+        }
+
+        #[inline]
+        fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
+            self.write_vectored(bufs)
+        }
+        #[inline]
+        fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
+            self.write_all(buf)
+        }
+
+        #[inline]
+        fn flush(&mut self) -> std::io::Result<()> {
+            self.flush()
+        }
+    }
+}

tests/test_buf.rs

@@ -118,3 +118,37 @@ fn copy_to_bytes_overflow() {
 
     let _bytes = buf.copy_to_bytes(12);
 }
+
+#[cfg(feature = "std")]
+#[test]
+fn fmt_write_bytesmut_test() {
+    use bytes::BytesMut;
+    use std::fmt::Write;
+    let mut buf = BytesMut::with_capacity(64);
+    write!(&mut buf, "hello").ok();
+    assert_eq!(&buf[..], b"hello");
+}
+
+#[cfg(feature = "std")]
+#[test]
+fn io_write_bytesmut_test() {
+    use bytes::ext::BytesMutExt;
+    use bytes::BytesMut;
+    use std::io::Write;
+    let mut buf = BytesMut::with_capacity(64);
+    write!(&mut buf as &mut dyn BytesMutExt, "hello").ok();
+    assert_eq!(&buf[..], b"hello");
+}
+
+#[cfg(feature = "std")]
+#[test]
+fn fmt_write_bytesmut_with_both_write_in_scope_test() {
+    use bytes::ext::BytesMutExt;
+    use bytes::BytesMut;
+    #[allow(unused_imports)]
+    use std::fmt::Write as FmtWrite;
+    use std::io::Write as IoWrite;
+    let mut buf = BytesMut::with_capacity(64);
+    write!(&mut buf as &mut dyn BytesMutExt, "hello").ok();
+    assert_eq!(&buf[..], b"hello");
+}