codec: allow AnyDelimiterCodec to take bytes

tokio-rs · Dec 13, 2023 · 624ac7f · 624ac7f
1 parent 4aa7bbf
commit 624ac7f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/tokio-util/src/codec/any_delimiter_codec.rs b/tokio-util/src/codec/any_delimiter_codec.rs
@@ -2,7 +2,7 @@ use crate::codec::decoder::Decoder;
 use crate::codec::encoder::Encoder;
 
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use std::{cmp, fmt, io, str, usize};
+use std::{cmp, fmt, io, usize};
 
 const DEFAULT_SEEK_DELIMITERS: &[u8] = b",;\n\r";
 const DEFAULT_SEQUENCE_WRITER: &[u8] = b",";
@@ -211,14 +211,14 @@ impl Decoder for AnyDelimiterCodec {
 
 impl<T> Encoder<T> for AnyDelimiterCodec
 where
-    T: AsRef<str>,
+    T: AsRef<[u8]>,
 {
     type Error = AnyDelimiterCodecError;
 
     fn encode(&mut self, chunk: T, buf: &mut BytesMut) -> Result<(), AnyDelimiterCodecError> {
         let chunk = chunk.as_ref();
         buf.reserve(chunk.len() + 1);
-        buf.put(chunk.as_bytes());
+        buf.put(chunk);
         buf.put(self.sequence_writer.as_ref());
 
         Ok(())

diff --git a/tokio-util/tests/codecs.rs b/tokio-util/tests/codecs.rs
@@ -462,3 +462,15 @@ fn any_delimiter_encoder() {
     codec.encode("chunk 2", &mut buf).unwrap();
     assert_eq!("chunk 1;--;chunk 2;--;", buf);
 }
+
+#[test]
+fn any_delimiter_encoder_bytes() {
+    let mut codec = AnyDelimiterCodec::new(b",".to_vec(), b";--;".to_vec());
+    let mut buf = BytesMut::new();
+
+    codec.encode(b"chunk 1", &mut buf).unwrap();
+    assert_eq!(b"chunk 1;--;", buf.as_ref());
+
+    codec.encode(b"chunk 2", &mut buf).unwrap();
+    assert_eq!(b"chunk 1;--;chunk 2;--;", buf.as_ref());
+}