Implement bytes.hex() with optional sep

qingshi163 · web-flow · commit 2d9b1d064940 · 2020-10-11T18:34:08.000+09:00
diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
@@ -424,8 +424,6 @@ def test_hex(self):
         self.assertEqual(self.type2test(b"\x1a\x2b\x30").hex(), '1a2b30')
         self.assertEqual(memoryview(b"\x1a\x2b\x30").hex(), '1a2b30')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_hex_separator_basics(self):
         three_bytes = self.type2test(b'\xb9\x01\xef')
         self.assertEqual(three_bytes.hex(), 'b901ef')
@@ -467,8 +465,6 @@ def test_hex_separator_five_bytes(self):
         five_bytes = self.type2test(range(90,95))
         self.assertEqual(five_bytes.hex(), '5a5b5c5d5e')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_hex_separator_six_bytes(self):
         six_bytes = self.type2test(x*3 for x in range(1, 7))
         self.assertEqual(six_bytes.hex(), '0306090c0f12')
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
@@ -3,27 +3,25 @@ use itertools::Itertools;
 use num_bigint::BigInt;
 use num_traits::ToPrimitive;
 
+use crate::anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper};
 use crate::byteslike::try_bytes_like;
 use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objbytearray::PyByteArray;
-use crate::obj::objbytes::PyBytes;
+use crate::obj::objbytes::{PyBytes, PyBytesRef};
 use crate::obj::objint::{self, PyInt, PyIntRef};
 use crate::obj::objlist::PyList;
 use crate::obj::objmemory::PyMemoryView;
 use crate::obj::objsingletons::PyNoneRef;
 use crate::obj::objslice::PySliceRef;
 use crate::obj::objstr::{self, PyStr, PyStrRef};
+use crate::obj::objtuple::PyTuple;
 use crate::pyobject::{
     BorrowValue, Either, PyComparisonValue, PyIterable, PyIterator, PyObjectRef, PyResult,
     TryFromObject, TypeProtocol,
 };
 use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
 use crate::slots::PyComparisonOp;
 use crate::vm::VirtualMachine;
-use crate::{
-    anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
-    obj::objtuple::PyTuple,
-};
 use rustpython_common::hash;
 
 #[derive(Debug, Default, Clone)]
@@ -549,11 +547,13 @@ impl PyBytesInner {
         new
     }
 
-    pub fn hex(&self) -> String {
-        self.elements
-            .iter()
-            .map(|x| format!("{:02x}", x))
-            .collect::<String>()
+    pub fn hex(
+        &self,
+        sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
+        bytes_per_sep: OptionalArg<isize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<String> {
+        bytes_to_hex(self.elements.as_slice(), sep, bytes_per_sep, vm)
     }
 
     pub fn fromhex(string: &str, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
@@ -1235,3 +1235,98 @@ pub fn bytes_decode(
             ))
         })
 }
+
+fn hex_impl_no_sep(bytes: &[u8]) -> String {
+    let mut buf: Vec<u8> = vec![0; bytes.len() * 2];
+    hex::encode_to_slice(bytes, buf.as_mut_slice()).unwrap();
+    unsafe { String::from_utf8_unchecked(buf) }
+}
+
+fn hex_impl(bytes: &[u8], sep: u8, bytes_per_sep: isize) -> String {
+    let len = bytes.len();
+
+    let buf = if bytes_per_sep < 0 {
+        let bytes_per_sep = std::cmp::min(len, (-bytes_per_sep) as usize);
+        let chunks = (len - 1) / bytes_per_sep;
+        let chunked = chunks * bytes_per_sep;
+        let unchunked = len - chunked;
+        let mut buf = vec![0; len * 2 + chunks];
+        let mut j = 0;
+        for i in (0..chunks).map(|i| i * bytes_per_sep) {
+            hex::encode_to_slice(
+                &bytes[i..i + bytes_per_sep],
+                &mut buf[j..j + bytes_per_sep * 2],
+            )
+            .unwrap();
+            j += bytes_per_sep * 2;
+            buf[j] = sep;
+            j += 1;
+        }
+        hex::encode_to_slice(&bytes[chunked..], &mut buf[j..j + unchunked * 2]).unwrap();
+        buf
+    } else {
+        let bytes_per_sep = std::cmp::min(len, bytes_per_sep as usize);
+        let chunks = (len - 1) / bytes_per_sep;
+        let chunked = chunks * bytes_per_sep;
+        let unchunked = len - chunked;
+        let mut buf = vec![0; len * 2 + chunks];
+        hex::encode_to_slice(&bytes[..unchunked], &mut buf[..unchunked * 2]).unwrap();
+        let mut j = unchunked * 2;
+        for i in (0..chunks).map(|i| i * bytes_per_sep + unchunked) {
+            buf[j] = sep;
+            j += 1;
+            hex::encode_to_slice(
+                &bytes[i..i + bytes_per_sep],
+                &mut buf[j..j + bytes_per_sep * 2],
+            )
+            .unwrap();
+            j += bytes_per_sep * 2;
+        }
+        buf
+    };
+
+    unsafe { String::from_utf8_unchecked(buf) }
+}
+
+pub fn bytes_to_hex(
+    bytes: &[u8],
+    sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
+    bytes_per_sep: OptionalArg<isize>,
+    vm: &VirtualMachine,
+) -> PyResult<String> {
+    if bytes.is_empty() {
+        return Ok("".to_owned());
+    }
+
+    if let OptionalArg::Present(sep) = sep {
+        let bytes_per_sep = bytes_per_sep.unwrap_or(1);
+        if bytes_per_sep == 0 {
+            return Ok(hex_impl_no_sep(bytes));
+        }
+
+        let s_guard;
+        let b_guard;
+        let sep = match &sep {
+            Either::A(s) => {
+                s_guard = s.borrow_value();
+                s_guard.as_bytes()
+            }
+            Either::B(bytes) => {
+                b_guard = bytes.borrow_value();
+                b_guard
+            }
+        };
+
+        if sep.len() != 1 {
+            return Err(vm.new_value_error("sep must be length 1.".to_owned()));
+        }
+        let sep = sep[0];
+        if sep > 127 {
+            return Err(vm.new_value_error("sep must be ASCII.".to_owned()));
+        }
+
+        Ok(hex_impl(bytes, sep, bytes_per_sep))
+    } else {
+        Ok(hex_impl_no_sep(bytes))
+    }
+}
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
@@ -1,33 +1,32 @@
 //! Implementation of the python bytearray object.
-use bstr::ByteSlice;
-use crossbeam_utils::atomic::AtomicCell;
-use rustpython_common::borrow::{BorrowedValue, BorrowedValueMut};
-use std::mem::size_of;
-
-use super::objint::PyIntRef;
-use super::objiter;
-use super::objstr::PyStrRef;
-use super::objtype::PyTypeRef;
 use crate::anystr::{self, AnyStr};
 use crate::bytesinner::{
     bytes_decode, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
     ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner,
 };
 use crate::byteslike::PyBytesLike;
+use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
 use crate::common::lock::{
     PyRwLock, PyRwLockReadGuard, PyRwLockUpgradableReadGuard, PyRwLockWriteGuard,
 };
 use crate::function::{OptionalArg, OptionalOption};
-use crate::obj::objbytes::PyBytes;
+use crate::obj::objbytes::{PyBytes, PyBytesRef};
+use crate::obj::objint::PyIntRef;
+use crate::obj::objiter;
 use crate::obj::objmemory::{Buffer, BufferOptions};
+use crate::obj::objstr::PyStrRef;
 use crate::obj::objtuple::PyTupleRef;
+use crate::obj::objtype::PyTypeRef;
 use crate::pyobject::{
     BorrowValue, Either, IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext,
     PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
 };
 use crate::sliceable::SequenceIndex;
 use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp, Unhashable};
 use crate::vm::VirtualMachine;
+use bstr::ByteSlice;
+use crossbeam_utils::atomic::AtomicCell;
+use std::mem::size_of;
 
 /// "bytearray(iterable_of_ints) -> bytearray\n\
 ///  bytearray(string, encoding[, errors]) -> bytearray\n\
@@ -249,8 +248,13 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "hex")]
-    fn hex(&self) -> String {
-        self.borrow_value().hex()
+    fn hex(
+        &self,
+        sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
+        bytes_per_sep: OptionalArg<isize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<String> {
+        self.borrow_value().hex(sep, bytes_per_sep, vm)
     }
 
     #[pymethod]
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -221,11 +221,14 @@ impl PyBytes {
         self.inner.swapcase().into()
     }
 
-    // TODO: Changed in version 3.8: bytes.hex() now supports optional sep and
-    // bytes_per_sep parameters to insert separators between bytes in the hex output.
     #[pymethod(name = "hex")]
-    pub(crate) fn hex(&self) -> String {
-        self.inner.hex()
+    pub(crate) fn hex(
+        &self,
+        sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
+        bytes_per_sep: OptionalArg<isize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<String> {
+        self.inner.hex(sep, bytes_per_sep, vm)
     }
 
     #[pymethod]
diff --git a/vm/src/obj/objmemory.rs b/vm/src/obj/objmemory.rs
@@ -1,15 +1,17 @@
 use std::{fmt::Debug, ops::Deref};
 
+use crate::bytesinner::bytes_to_hex;
 use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
 use crate::common::hash::PyHash;
+use crate::function::OptionalArg;
 use crate::obj::objbytes::{PyBytes, PyBytesRef};
 use crate::obj::objlist::{PyList, PyListRef};
 use crate::obj::objslice::PySliceRef;
-use crate::obj::objstr::PyStr;
+use crate::obj::objstr::{PyStr, PyStrRef};
 use crate::obj::objtype::PyTypeRef;
 use crate::pyobject::{
-    IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef,
-    PyResult, PyThreadingConstraint, PyValue, TypeProtocol,
+    Either, IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef,
+    PyRef, PyResult, PyThreadingConstraint, PyValue, TypeProtocol,
 };
 use crate::sliceable::{convert_slice, saturate_range, wrap_index, SequenceIndex};
 use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp};
@@ -592,10 +594,13 @@ impl PyMemoryView {
         }
     }
 
-    // TODO: Changed in version 3.8: memoryview.hex() now supports optional sep and bytes_per_sep
-    // parameters to insert separators between bytes in the hex output.
     #[pymethod]
-    fn hex(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<String> {
+    fn hex(
+        zelf: PyRef<Self>,
+        sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
+        bytes_per_sep: OptionalArg<isize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<String> {
         zelf.try_not_released(vm)?;
         let guard;
         let vec;
@@ -609,11 +614,8 @@ impl PyMemoryView {
                 vec.as_slice()
             }
         };
-        let s = bytes
-            .iter()
-            .map(|x| format!("{:02x}", x))
-            .collect::<String>();
-        Ok(s)
+
+        bytes_to_hex(bytes, sep, bytes_per_sep, vm)
     }
 
     fn eq(zelf: &PyRef<Self>, other: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
