Skip to content

Commit e1ddbda

Browse files
qingshi163qingshi163
committed
fix unittest test_array.test_buffer
1 parent 4489ba1 commit e1ddbda

File tree

10 files changed

+275
-217
lines changed

10 files changed

+275
-217
lines changed

‎Lib/test/test_array.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1013,8 +1013,6 @@ def test_coveritertraverse(self):
10131013
l.append(l)
10141014
gc.collect()
10151015

1016-
# TODO: RUSTPYTHON
1017-
@unittest.expectedFailure
10181016
def test_buffer(self):
10191017
a = array.array(self.typecode, self.example)
10201018
m = memoryview(a)

‎vm/src/bytesinner.rs

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ use itertools::Itertools;
33
use num_bigint::BigInt;
44
use num_traits::ToPrimitive;
55

6-
use crate::anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper};
76
use crate::byteslike::{try_bytes_like, PyBytesLike};
87
use crate::function::{OptionalArg, OptionalOption};
98
use crate::obj::objbytearray::PyByteArray;
@@ -21,6 +20,10 @@ use crate::pyobject::{
2120
use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
2221
use crate::slots::PyComparisonOp;
2322
use crate::vm::VirtualMachine;
23+
use crate::{
24+
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
25+
obj::objmemory::try_buffer_from_object,
26+
};
2427
use rustpython_common::hash;
2528

2629
#[derive(Debug, Default, Clone)]
@@ -36,17 +39,16 @@ impl From<Vec<u8>> for PyBytesInner {
3639

3740
impl TryFromObject for PyBytesInner {
3841
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
42+
if let Ok(buffer) = try_buffer_from_object(obj.clone(), vm) {
43+
let bytes = buffer
44+
.as_contiguous()
45+
.ok_or_else(|| vm.new_buffer_error("buffer is not contiguous".to_owned()))?;
46+
return Ok(Self::from(bytes.to_vec()));
47+
}
48+
3949
match_class!(match obj {
40-
i @ PyBytes => Ok(PyBytesInner {
41-
elements: i.borrow_value().to_vec()
42-
}),
43-
j @ PyByteArray => Ok(PyBytesInner {
44-
elements: j.borrow_value().elements.to_vec()
45-
}),
46-
k @ PyMemoryView => Ok(PyBytesInner {
47-
elements: k.try_bytes(|v| v.to_vec()).unwrap()
48-
}),
4950
l @ PyList => l.to_byte_inner(vm),
51+
// TODO: PyTyple
5052
obj => {
5153
let iter = vm.get_method_or_type_error(obj.clone(), "__iter__", || {
5254
format!(
@@ -344,7 +346,7 @@ impl PyBytesInner {
344346
iter.map(|obj| Self::value_try_from_object(vm, obj?))
345347
.try_collect()
346348
} else if let Some(mview) = object.payload_if_subclass::<PyMemoryView>(vm) {
347-
Ok(mview.try_bytes(|v| v.to_vec()).unwrap())
349+
mview.try_bytes(vm, |v| v.to_vec())
348350
} else {
349351
Err(vm.new_type_error(
350352
"can assign only bytes, buffers, or iterables of ints in range(0, 256)".to_owned(),
@@ -362,11 +364,32 @@ impl PyBytesInner {
362364
self.elements.set_slice_items(vm, &slice, items.as_slice())
363365
}
364366

367+
pub fn setslice_no_resize(
368+
&mut self,
369+
slice: PySliceRef,
370+
object: PyObjectRef,
371+
vm: &VirtualMachine,
372+
) -> PyResult<()> {
373+
let items = Self::value_seq_try_from_object(vm, object)?;
374+
self.elements
375+
.set_slice_items_no_resize(vm, &slice, items.as_slice())
376+
}
377+
365378
pub fn setslice_from_self(&mut self, slice: PySliceRef, vm: &VirtualMachine) -> PyResult<()> {
366379
let items = self.elements.clone();
367380
self.elements.set_slice_items(vm, &slice, items.as_slice())
368381
}
369382

383+
pub fn setslice_from_self_no_resize(
384+
&mut self,
385+
slice: PySliceRef,
386+
vm: &VirtualMachine,
387+
) -> PyResult<()> {
388+
let items = self.elements.clone();
389+
self.elements
390+
.set_slice_items_no_resize(vm, &slice, items.as_slice())
391+
}
392+
370393
pub fn delitem(&mut self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult<()> {
371394
match needle {
372395
SequenceIndex::Int(int) => {

‎vm/src/obj/objbytearray.rs

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,7 @@ use crate::pyobject::{
2626
PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
2727
};
2828
use crate::sliceable::SequenceIndex;
29-
use crate::slots::BufferProtocol;
30-
use crate::slots::{Comparable, Hashable, PyComparisonOp, Unhashable};
29+
use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp, Unhashable};
3130
use crate::vm::VirtualMachine;
3231

3332
/// "bytearray(iterable_of_ints) -> bytearray\n\
@@ -173,13 +172,14 @@ impl PyByteArray {
173172
) -> PyResult<()> {
174173
match needle {
175174
SequenceIndex::Int(int) => zelf.borrow_value_mut().setindex(int, value, vm),
176-
SequenceIndex::Slice(slice) => {
177-
if zelf.is(&value) {
178-
zelf.borrow_value_mut().setslice_from_self(slice, vm)
179-
} else {
180-
zelf.borrow_value_mut().setslice(slice, value, vm)
181-
}
182-
}
175+
SequenceIndex::Slice(slice) => match (zelf.is(&value), zelf.is_resizable()) {
176+
(true, true) => zelf.borrow_value_mut().setslice_from_self(slice, vm),
177+
(true, false) => zelf
178+
.borrow_value_mut()
179+
.setslice_from_self_no_resize(slice, vm),
180+
(false, true) => zelf.borrow_value_mut().setslice(slice, value, vm),
181+
(false, false) => zelf.borrow_value_mut().setslice_no_resize(slice, value, vm),
182+
},
183183
}
184184
}
185185

‎vm/src/obj/objbytes.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,15 @@ impl From<Vec<u8>> for PyBytes {
6464
}
6565
}
6666

67+
impl From<PyBytesInner> for PyBytes {
68+
fn from(inner: PyBytesInner) -> Self {
69+
Self {
70+
inner,
71+
buffer_options: OnceCell::new(),
72+
}
73+
}
74+
}
75+
6776
impl IntoPyObject for Vec<u8> {
6877
fn into_pyobject(self, vm: &VirtualMachine) -> PyObjectRef {
6978
vm.ctx.new_bytes(self)
@@ -549,3 +558,9 @@ impl PyBytesIterator {
549558
zelf
550559
}
551560
}
561+
562+
impl TryFromObject for PyBytes {
563+
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
564+
PyBytesInner::try_from_object(vm, obj).map(|x| x.into())
565+
}
566+
}

‎vm/src/obj/objint.rs

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use super::objbool::IntoPyBool;
1111
use super::objbytearray::PyByteArray;
1212
use super::objbytes::PyBytes;
1313
use super::objfloat;
14-
use super::objmemory::PyMemoryView;
14+
use super::objmemory::try_buffer_from_object;
1515
use super::objstr::{PyStr, PyStrRef};
1616
use super::objtype::PyTypeRef;
1717
use crate::bytesinner::PyBytesInner;
@@ -733,31 +733,31 @@ pub(crate) fn to_int(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<BigInt>
733733
let inner = bytearray.borrow_value();
734734
bytes_to_int(&inner.elements, base)
735735
}
736-
memoryview @ PyMemoryView => {
737-
// TODO: proper error handling instead of `unwrap()`
738-
memoryview
739-
.try_bytes(|bytes| bytes_to_int(&bytes, base))
740-
.unwrap()
741-
}
742736
array @ PyArray => {
743-
let bytes = array.tobytes();
744-
bytes_to_int(&bytes, base)
737+
let bytes = array.get_bytes();
738+
bytes_to_int(&*bytes, base)
745739
}
746740
obj => {
747-
let method = vm.get_method_or_type_error(obj.clone(), "__int__", || {
748-
format!(
749-
"int() argument must be a string or a number, not '{}'",
750-
obj.lease_class().name
751-
)
752-
})?;
753-
let result = vm.invoke(&method, PyFuncArgs::default())?;
754-
return match result.payload::<PyInt>() {
755-
Some(int_obj) => Ok(int_obj.borrow_value().clone()),
756-
None => Err(vm.new_type_error(format!(
757-
"TypeError: __int__ returned non-int (type '{}')",
758-
result.lease_class().name
759-
))),
760-
};
741+
if let Some(method) = vm.get_method(obj.clone(), "__int__") {
742+
let result = vm.invoke(&method?, PyFuncArgs::default())?;
743+
return match result.payload::<PyInt>() {
744+
Some(int_obj) => Ok(int_obj.borrow_value().clone()),
745+
None => Err(vm.new_type_error(format!(
746+
"TypeError: __int__ returned non-int (type '{}')",
747+
result.class().name
748+
))),
749+
};
750+
}
751+
752+
if let Ok(bytes) = try_buffer_from_object(obj.clone(), vm) {
753+
let bytes = bytes.as_contiguous()
754+
.ok_or_else(|| vm.new_value_error(
755+
format!("int() argument must be a string, a bytes-like object or a number, not '{}'",
756+
obj.lease_class().name)))?;
757+
bytes_to_int(&*bytes, base)
758+
} else {
759+
None
760+
}
761761
}
762762
});
763763
match opt {

‎vm/src/obj/objmemory.rs

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,9 @@ use crate::pyobject::{
99
PyValue, TypeProtocol,
1010
};
1111
use crate::sliceable::{saturate_range, wrap_index, SequenceIndex};
12-
use crate::slots::{BufferProtocol, Hashable};
12+
use crate::slots::{BufferProtocol, Comparable, Hashable};
1313
use crate::stdlib::pystruct::_struct::FormatSpec;
1414
use crate::VirtualMachine;
15-
use crate::{bytesinner::try_as_bytes, slots::Comparable};
1615
use crate::{common::hash::PyHash, pyobject::PyComparisonValue};
1716
use crossbeam_utils::atomic::AtomicCell;
1817
use itertools::Itertools;
@@ -140,11 +139,14 @@ impl PyMemoryView {
140139
})
141140
}
142141

143-
pub fn try_bytes<F, R>(&self, f: F) -> Option<R>
142+
pub fn try_bytes<F, R>(&self, vm: &VirtualMachine, f: F) -> PyResult<R>
144143
where
145-
F: Fn(&[u8]) -> R,
144+
F: FnOnce(&[u8]) -> R,
146145
{
147-
try_as_bytes(self.obj.clone(), f)
146+
self.try_not_released(vm)?;
147+
self.buffer.as_contiguous().map(|x| f(&*x)).ok_or_else(|| {
148+
vm.new_type_error("non-contiguous memoryview is not a bytes-like object".to_owned())
149+
})
148150
}
149151

150152
#[pyslot]
@@ -562,7 +564,7 @@ pub fn try_buffer_from_object(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult
562564
.and_then(|buffer_func| buffer_func(obj, vm).ok().map(|x| BufferRef(x)))
563565
.ok_or_else(|| {
564566
vm.new_type_error(format!(
565-
"memoryview: a bytes-like object is required, not '{}'",
567+
"a bytes-like object is required, not '{}'",
566568
obj_cls.name
567569
))
568570
})

0 commit comments

Comments
 (0)