Skip to content

Commit 03a6798

Browse files
youknowoneyouknowone
authored
Merge pull request RustPython#3243 from youknowone/cleanup-function-byteslike
Cleanup function byteslike
2 parents 5bee8fc + 04e25e8 commit 03a6798

3 files changed

Lines changed: 73 additions & 80 deletions

File tree

‎vm/src/function.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
mod argument;
2-
mod byteslike;
2+
mod buffer;
33
mod number;
44

55
use crate::{
@@ -14,7 +14,7 @@ use std::marker::PhantomData;
1414
use std::ops::RangeInclusive;
1515

1616
pub use argument::{ArgCallable, ArgIterable};
17-
pub use byteslike::{ArgBytesLike, ArgMemoryBuffer, ArgStrOrBytesLike};
17+
pub use buffer::{ArgBytesLike, ArgMemoryBuffer, ArgStrOrBytesLike};
1818
pub use number::{ArgIntoBool, ArgIntoComplex, ArgIntoFloat};
1919

2020
/// Implemented by any type that can be returned from a built-in Python function.

vm/src/function/byteslike.rs renamed to vm/src/function/buffer.rs

Lines changed: 64 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,48 @@
1-
use crate::builtins::PyStrRef;
2-
use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
3-
use crate::protocol::PyBuffer;
4-
use crate::vm::VirtualMachine;
5-
use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject};
1+
use crate::{
2+
builtins::PyStrRef,
3+
common::borrow::{BorrowedValue, BorrowedValueMut},
4+
protocol::PyBuffer,
5+
PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject, VirtualMachine,
6+
};
67

78
// Python/getargs.c
89

910
/// any bytes-like object. Like the `y*` format code for `PyArg_Parse` in CPython.
1011
#[derive(Debug)]
1112
pub struct ArgBytesLike(PyBuffer);
1213

13-
/// A memory buffer, read-write access. Like the `w*` format code for `PyArg_Parse` in CPython.
14-
#[derive(Debug)]
15-
pub struct ArgMemoryBuffer(PyBuffer);
14+
impl PyObjectRef {
15+
pub fn try_bytes_like<R>(
16+
&self,
17+
vm: &VirtualMachine,
18+
f: impl FnOnce(&[u8]) -> R,
19+
) -> PyResult<R> {
20+
let buffer = PyBuffer::try_from_borrowed_object(vm, self)?;
21+
buffer.as_contiguous().map(|x| f(&*x)).ok_or_else(|| {
22+
vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())
23+
})
24+
}
25+
26+
pub fn try_rw_bytes_like<R>(
27+
&self,
28+
vm: &VirtualMachine,
29+
f: impl FnOnce(&mut [u8]) -> R,
30+
) -> PyResult<R> {
31+
let buffer = PyBuffer::try_from_borrowed_object(vm, self)?;
32+
buffer
33+
.as_contiguous_mut()
34+
.map(|mut x| f(&mut *x))
35+
.ok_or_else(|| {
36+
vm.new_type_error("buffer is not a read-write bytes-like object".to_owned())
37+
})
38+
}
39+
}
1640

1741
impl ArgBytesLike {
42+
pub fn borrow_buf(&self) -> BorrowedValue<'_, [u8]> {
43+
self.0.as_contiguous().unwrap()
44+
}
45+
1846
pub fn with_ref<F, R>(&self, f: F) -> R
1947
where
2048
F: FnOnce(&[u8]) -> R,
@@ -35,77 +63,56 @@ impl ArgBytesLike {
3563
}
3664
}
3765

38-
impl ArgMemoryBuffer {
39-
pub fn with_ref<F, R>(&self, f: F) -> R
40-
where
41-
F: FnOnce(&mut [u8]) -> R,
42-
{
43-
f(&mut *self.borrow_buf_mut())
44-
}
45-
46-
pub fn len(&self) -> usize {
47-
self.borrow_buf_mut().len()
48-
}
49-
50-
pub fn is_empty(&self) -> bool {
51-
self.borrow_buf_mut().is_empty()
66+
impl From<ArgBytesLike> for PyBuffer {
67+
fn from(buffer: ArgBytesLike) -> Self {
68+
buffer.0
5269
}
5370
}
5471

55-
impl ArgBytesLike {
56-
pub fn new(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
72+
impl TryFromBorrowedObject for ArgBytesLike {
73+
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
5774
let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?;
5875
if buffer.options.contiguous {
5976
Ok(Self(buffer))
6077
} else {
6178
Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()))
6279
}
6380
}
81+
}
82+
83+
/// A memory buffer, read-write access. Like the `w*` format code for `PyArg_Parse` in CPython.
84+
#[derive(Debug)]
85+
pub struct ArgMemoryBuffer(PyBuffer);
6486

65-
pub fn into_buffer(self) -> PyBuffer {
66-
self.0
87+
impl ArgMemoryBuffer {
88+
pub fn borrow_buf_mut(&self) -> BorrowedValueMut<'_, [u8]> {
89+
self.0.as_contiguous_mut().unwrap()
6790
}
6891

69-
pub fn borrow_buf(&self) -> BorrowedValue<'_, [u8]> {
70-
self.0.as_contiguous().unwrap()
92+
pub fn with_ref<F, R>(&self, f: F) -> R
93+
where
94+
F: FnOnce(&mut [u8]) -> R,
95+
{
96+
f(&mut *self.borrow_buf_mut())
7197
}
72-
}
7398

74-
impl TryFromBorrowedObject for ArgBytesLike {
75-
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
76-
Self::new(vm, obj)
99+
pub fn len(&self) -> usize {
100+
self.borrow_buf_mut().len()
77101
}
78-
}
79102

80-
impl PyObjectRef {
81-
pub fn try_bytes_like<R>(
82-
&self,
83-
vm: &VirtualMachine,
84-
f: impl FnOnce(&[u8]) -> R,
85-
) -> PyResult<R> {
86-
let buffer = PyBuffer::try_from_borrowed_object(vm, self)?;
87-
buffer.as_contiguous().map(|x| f(&*x)).ok_or_else(|| {
88-
vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())
89-
})
103+
pub fn is_empty(&self) -> bool {
104+
self.borrow_buf_mut().is_empty()
90105
}
106+
}
91107

92-
pub fn try_rw_bytes_like<R>(
93-
&self,
94-
vm: &VirtualMachine,
95-
f: impl FnOnce(&mut [u8]) -> R,
96-
) -> PyResult<R> {
97-
let buffer = PyBuffer::try_from_borrowed_object(vm, self)?;
98-
buffer
99-
.as_contiguous_mut()
100-
.map(|mut x| f(&mut *x))
101-
.ok_or_else(|| {
102-
vm.new_type_error("buffer is not a read-write bytes-like object".to_owned())
103-
})
108+
impl From<ArgMemoryBuffer> for PyBuffer {
109+
fn from(buffer: ArgMemoryBuffer) -> Self {
110+
buffer.0
104111
}
105112
}
106113

107-
impl ArgMemoryBuffer {
108-
pub fn new(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
114+
impl TryFromBorrowedObject for ArgMemoryBuffer {
115+
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
109116
let buffer = PyBuffer::try_from_borrowed_object(vm, obj)?;
110117
if !buffer.options.contiguous {
111118
Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()))
@@ -115,20 +122,6 @@ impl ArgMemoryBuffer {
115122
Ok(Self(buffer))
116123
}
117124
}
118-
119-
pub fn into_buffer(self) -> PyBuffer {
120-
self.0
121-
}
122-
123-
pub fn borrow_buf_mut(&self) -> BorrowedValueMut<'_, [u8]> {
124-
self.0.as_contiguous_mut().unwrap()
125-
}
126-
}
127-
128-
impl TryFromBorrowedObject for ArgMemoryBuffer {
129-
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
130-
Self::new(vm, obj)
131-
}
132125
}
133126

134127
/// A text string or bytes-like object. Like the `s*` format code for `PyArg_Parse` in CPython.

‎vm/src/stdlib/io.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@ mod _io {
9090
slots::{Iterable, SlotConstructor, SlotDestructor, SlotIterator},
9191
utils::Either,
9292
vm::{ReprGuard, VirtualMachine},
93-
IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject,
94-
TypeProtocol,
93+
IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, StaticType,
94+
TryFromBorrowedObject, TryFromObject, TypeProtocol,
9595
};
9696
use bstr::ByteSlice;
9797
use crossbeam_utils::atomic::AtomicCell;
@@ -627,7 +627,7 @@ mod _io {
627627
method: &str,
628628
vm: &VirtualMachine,
629629
) -> PyResult<usize> {
630-
let b = ArgMemoryBuffer::new(vm, &bufobj)?;
630+
let b = ArgMemoryBuffer::try_from_borrowed_object(vm, &bufobj)?;
631631
let l = b.len();
632632
let data = vm.call_method(&zelf, method, (l,))?;
633633
if data.is(&bufobj) {
@@ -957,7 +957,7 @@ mod _io {
957957

958958
let mut remaining = buf_len;
959959
let mut written = 0;
960-
let buffer = obj.into_buffer();
960+
let buffer: PyBuffer = obj.into();
961961
while remaining > self.buffer.len() {
962962
let res = self.raw_write(Some(buffer.clone()), written..buf_len, vm)?;
963963
match res {
@@ -1633,14 +1633,14 @@ mod _io {
16331633
let mut data = self.reader().lock(vm)?;
16341634
let raw = data.check_init(vm)?;
16351635
ensure_unclosed(raw, "readinto of closed file", vm)?;
1636-
data.readinto_generic(buf.into_buffer(), false, vm)
1636+
data.readinto_generic(buf.into(), false, vm)
16371637
}
16381638
#[pymethod]
16391639
fn readinto1(&self, buf: ArgMemoryBuffer, vm: &VirtualMachine) -> PyResult<Option<usize>> {
16401640
let mut data = self.reader().lock(vm)?;
16411641
let raw = data.check_init(vm)?;
16421642
ensure_unclosed(raw, "readinto of closed file", vm)?;
1643-
data.readinto_generic(buf.into_buffer(), true, vm)
1643+
data.readinto_generic(buf.into(), true, vm)
16441644
}
16451645
}
16461646

@@ -2925,7 +2925,7 @@ mod _io {
29252925
let chunk_size = std::cmp::max(self.chunk_size, size_hint);
29262926
let input_chunk = vm.call_method(&self.buffer, method, (chunk_size,))?;
29272927

2928-
let buf = ArgBytesLike::new(vm, &input_chunk).map_err(|_| {
2928+
let buf = ArgBytesLike::try_from_borrowed_object(vm, &input_chunk).map_err(|_| {
29292929
vm.new_type_error(format!(
29302930
"underlying {}() should have returned a bytes-like object, not '{}'",
29312931
method,

0 commit comments

Comments
 (0)