Merge pull request #3243 from youknowone/cleanup-function-byteslike

Cleanup function byteslike

vm/src/function.rs

@@ -1,5 +1,5 @@
 mod argument;
-mod byteslike;
+mod buffer;
 mod number;
 
 use crate::{
@@ -14,7 +14,7 @@ use std::marker::PhantomData;
 use std::ops::RangeInclusive;
 
 pub use argument::{ArgCallable, ArgIterable};
-pub use byteslike::{ArgBytesLike, ArgMemoryBuffer, ArgStrOrBytesLike};
+pub use buffer::{ArgBytesLike, ArgMemoryBuffer, ArgStrOrBytesLike};
 pub use number::{ArgIntoBool, ArgIntoComplex, ArgIntoFloat};
 
 /// Implemented by any type that can be returned from a built-in Python function.

vm/src/function/buffer.rs

@@ -1,8 +1,9 @@
-use crate::builtins::PyStrRef;
-use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
-use crate::protocol::PyBuffer;
-use crate::vm::VirtualMachine;
-use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject};
+use crate::{
+    builtins::PyStrRef,
+    common::borrow::{BorrowedValue, BorrowedValueMut},
+    protocol::PyBuffer,
+    PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject, VirtualMachine,
+};
 
 // Python/getargs.c
 
@@ -10,73 +11,6 @@ use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject};
 #[derive(Debug)]
 pub struct ArgBytesLike(PyBuffer);
 
-/// A memory buffer, read-write access. Like the `w*` format code for `PyArg_Parse` in CPython.
-#[derive(Debug)]
-pub struct ArgMemoryBuffer(PyBuffer);
-
-impl ArgBytesLike {
-    pub fn with_ref<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(&[u8]) -> R,
-    {
-        f(&*self.borrow_buf())
-    }
-
-    pub fn len(&self) -> usize {
-        self.borrow_buf().len()
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.borrow_buf().is_empty()
-    }
-
-    pub fn to_cow(&self) -> std::borrow::Cow<[u8]> {
-        self.borrow_buf().to_vec().into()
-    }
-}
-
-impl ArgMemoryBuffer {
-    pub fn with_ref<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(&mut [u8]) -> R,
-    {
-        f(&mut *self.borrow_buf_mut())
-    }
-
-    pub fn len(&self) -> usize {
-        self.borrow_buf_mut().len()
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.borrow_buf_mut().is_empty()
-    }
-}
-
-impl ArgBytesLike {
-    pub fn new(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
-        let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?;
-        if buffer.options.contiguous {
-            Ok(Self(buffer))
-        } else {
-            Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()))
-        }
-    }
-
-    pub fn into_buffer(self) -> PyBuffer {
-        self.0
-    }
-
-    pub fn borrow_buf(&self) -> BorrowedValue<'_, [u8]> {
-        self.0.as_contiguous().unwrap()
-    }
-}
-
-impl TryFromBorrowedObject for ArgBytesLike {
-    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
-        Self::new(vm, obj)
-    }
-}
-
 impl PyObjectRef {
     pub fn try_bytes_like<R>(
         &self,
@@ -104,8 +38,81 @@ impl PyObjectRef {
     }
 }
 
+impl ArgBytesLike {
+    pub fn borrow_buf(&self) -> BorrowedValue<'_, [u8]> {
+        self.0.as_contiguous().unwrap()
+    }
+
+    pub fn with_ref<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&[u8]) -> R,
+    {
+        f(&*self.borrow_buf())
+    }
+
+    pub fn len(&self) -> usize {
+        self.borrow_buf().len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.borrow_buf().is_empty()
+    }
+
+    pub fn to_cow(&self) -> std::borrow::Cow<[u8]> {
+        self.borrow_buf().to_vec().into()
+    }
+}
+
+impl From<ArgBytesLike> for PyBuffer {
+    fn from(buffer: ArgBytesLike) -> Self {
+        buffer.0
+    }
+}
+
+impl TryFromBorrowedObject for ArgBytesLike {
+    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
+        let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?;
+        if buffer.options.contiguous {
+            Ok(Self(buffer))
+        } else {
+            Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()))
+        }
+    }
+}
+
+/// A memory buffer, read-write access. Like the `w*` format code for `PyArg_Parse` in CPython.
+#[derive(Debug)]
+pub struct ArgMemoryBuffer(PyBuffer);
+
 impl ArgMemoryBuffer {
-    pub fn new(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
+    pub fn borrow_buf_mut(&self) -> BorrowedValueMut<'_, [u8]> {
+        self.0.as_contiguous_mut().unwrap()
+    }
+
+    pub fn with_ref<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R,
+    {
+        f(&mut *self.borrow_buf_mut())
+    }
+
+    pub fn len(&self) -> usize {
+        self.borrow_buf_mut().len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.borrow_buf_mut().is_empty()
+    }
+}
+
+impl From<ArgMemoryBuffer> for PyBuffer {
+    fn from(buffer: ArgMemoryBuffer) -> Self {
+        buffer.0
+    }
+}
+
+impl TryFromBorrowedObject for ArgMemoryBuffer {
+    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
         let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?;
         if !buffer.options.contiguous {
             Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()))
@@ -115,20 +122,6 @@ impl ArgMemoryBuffer {
             Ok(Self(buffer))
         }
     }
-
-    pub fn into_buffer(self) -> PyBuffer {
-        self.0
-    }
-
-    pub fn borrow_buf_mut(&self) -> BorrowedValueMut<'_, [u8]> {
-        self.0.as_contiguous_mut().unwrap()
-    }
-}
-
-impl TryFromBorrowedObject for ArgMemoryBuffer {
-    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
-        Self::new(vm, obj)
-    }
 }
 
 /// A text string or bytes-like object. Like the `s*` format code for `PyArg_Parse` in CPython.

vm/src/stdlib/io.rs

@@ -90,8 +90,8 @@ mod _io {
         slots::{Iterable, SlotConstructor, SlotDestructor, SlotIterator},
         utils::Either,
         vm::{ReprGuard, VirtualMachine},
-        IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject,
-        TypeProtocol,
+        IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, StaticType,
+        TryFromBorrowedObject, TryFromObject, TypeProtocol,
     };
     use bstr::ByteSlice;
     use crossbeam_utils::atomic::AtomicCell;
@@ -627,7 +627,7 @@ mod _io {
             method: &str,
             vm: &VirtualMachine,
         ) -> PyResult<usize> {
-            let b = ArgMemoryBuffer::new(vm, &bufobj)?;
+            let b = ArgMemoryBuffer::try_from_borrowed_object(vm, &bufobj)?;
             let l = b.len();
             let data = vm.call_method(&zelf, method, (l,))?;
             if data.is(&bufobj) {
@@ -957,7 +957,7 @@ mod _io {
 
             let mut remaining = buf_len;
             let mut written = 0;
-            let buffer = obj.into_buffer();
+            let buffer: PyBuffer = obj.into();
             while remaining > self.buffer.len() {
                 let res = self.raw_write(Some(buffer.clone()), written..buf_len, vm)?;
                 match res {
@@ -1633,14 +1633,14 @@ mod _io {
             let mut data = self.reader().lock(vm)?;
             let raw = data.check_init(vm)?;
             ensure_unclosed(raw, "readinto of closed file", vm)?;
-            data.readinto_generic(buf.into_buffer(), false, vm)
+            data.readinto_generic(buf.into(), false, vm)
         }
         #[pymethod]
         fn readinto1(&self, buf: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<Option<usize>> {
             let mut data = self.reader().lock(vm)?;
             let raw = data.check_init(vm)?;
             ensure_unclosed(raw, "readinto of closed file", vm)?;
-            data.readinto_generic(buf.into_buffer(), true, vm)
+            data.readinto_generic(buf.into(), true, vm)
         }
     }
 
@@ -2925,7 +2925,7 @@ mod _io {
             let chunk_size = std::cmp::max(self.chunk_size, size_hint);
             let input_chunk = vm.call_method(&self.buffer, method, (chunk_size,))?;
 
-            let buf = ArgBytesLike::new(vm, &input_chunk).map_err(|_| {
+            let buf = ArgBytesLike::try_from_borrowed_object(vm, &input_chunk).map_err(|_| {
                 vm.new_type_error(format!(
                     "underlying {}() should have returned a bytes-like object, not '{}'",
                     method,