Skip to content

Commit 944eb0d

Browse files
qingshi163qingshi163
committed
Impl sequence protocol for bytearray
1 parent 4692dd1 commit 944eb0d

File tree

5 files changed

+85
-44
lines changed

5 files changed

+85
-44
lines changed

‎vm/src/builtins/bytearray.rs

Lines changed: 75 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crate::common::{
99
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard,
1010
PyRwLockWriteGuard,
1111
},
12+
static_cell,
1213
};
1314
use crate::{
1415
anystr::{self, AnyStr},
@@ -21,19 +22,20 @@ use crate::{
2122
function::{ArgBytesLike, ArgIterable, FuncArgs, IntoPyObject, OptionalArg, OptionalOption},
2223
protocol::{
2324
BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
25+
PySequenceMethods,
2426
},
2527
sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex},
2628
types::{
27-
AsBuffer, AsMapping, Callable, Comparable, Constructor, Hashable, IterNext,
29+
AsBuffer, AsMapping, AsSequence, Callable, Comparable, Constructor, Hashable, IterNext,
2830
IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
2931
},
3032
utils::Either,
3133
IdProtocol, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef,
32-
PyResult, PyValue, TypeProtocol, VirtualMachine,
34+
PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
3335
};
3436
use bstr::ByteSlice;
3537
use crossbeam_utils::atomic::AtomicCell;
36-
use std::mem::size_of;
38+
use std::{borrow::Cow, mem::size_of};
3739

3840
/// "bytearray(iterable_of_ints) -> bytearray\n\
3941
/// bytearray(string, encoding[, errors]) -> bytearray\n\
@@ -102,7 +104,7 @@ pub(crate) fn init(context: &PyContext) {
102104

103105
#[pyimpl(
104106
flags(BASETYPE),
105-
with(Hashable, Comparable, AsBuffer, AsMapping, Iterable)
107+
with(Hashable, Comparable, AsBuffer, AsMapping, AsSequence, Iterable)
106108
)]
107109
impl PyByteArray {
108110
#[pyslot]
@@ -164,6 +166,17 @@ impl PyByteArray {
164166
self.inner().contains(needle, vm)
165167
}
166168

169+
fn setitem_by_idx(&self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
170+
let value = value_from_object(vm, &value)?;
171+
let mut elements = self.borrow_buf_mut();
172+
if let Some(i) = elements.wrap_index(i) {
173+
elements[i] = value;
174+
Ok(())
175+
} else {
176+
Err(vm.new_index_error("index out of range".to_owned()))
177+
}
178+
}
179+
167180
#[pymethod(magic)]
168181
fn setitem(
169182
zelf: PyRef<Self>,
@@ -172,16 +185,7 @@ impl PyByteArray {
172185
vm: &VirtualMachine,
173186
) -> PyResult<()> {
174187
match SequenceIndex::try_from_object_for(vm, needle, Self::NAME)? {
175-
SequenceIndex::Int(i) => {
176-
let value = value_from_object(vm, &value)?;
177-
let mut elements = zelf.borrow_buf_mut();
178-
if let Some(i) = elements.wrap_index(i) {
179-
elements[i] = value;
180-
Ok(())
181-
} else {
182-
Err(vm.new_index_error("index out of range".to_owned()))
183-
}
184-
}
188+
SequenceIndex::Int(i) => zelf.setitem_by_idx(i, value, vm),
185189
SequenceIndex::Slice(slice) => {
186190
let slice = slice.to_saturated(vm)?;
187191
let items = if zelf.is(&value) {
@@ -212,18 +216,20 @@ impl PyByteArray {
212216
self.inner().getitem(Self::NAME, needle, vm)
213217
}
214218

219+
fn delitem_by_idx(&self, i: isize, vm: &VirtualMachine) -> PyResult<()> {
220+
let elements = &mut self.try_resizable(vm)?.elements;
221+
if let Some(idx) = elements.wrap_index(i) {
222+
elements.remove(idx);
223+
Ok(())
224+
} else {
225+
Err(vm.new_index_error("index out of range".to_owned()))
226+
}
227+
}
228+
215229
#[pymethod(magic)]
216230
pub fn delitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
217231
match SequenceIndex::try_from_object_for(vm, needle, Self::NAME)? {
218-
SequenceIndex::Int(int) => {
219-
let elements = &mut self.try_resizable(vm)?.elements;
220-
if let Some(idx) = elements.wrap_index(int) {
221-
elements.remove(idx);
222-
Ok(())
223-
} else {
224-
Err(vm.new_index_error("index out of range".to_owned()))
225-
}
226-
}
232+
SequenceIndex::Int(i) => self.delitem_by_idx(i, vm),
227233
SequenceIndex::Slice(slice) => {
228234
let slice = slice.to_saturated(vm)?;
229235
let elements = &mut self.try_resizable(vm)?.elements;
@@ -790,6 +796,52 @@ impl AsMapping for PyByteArray {
790796
}
791797
}
792798

799+
impl AsSequence for PyByteArray {
800+
fn as_sequence(_zelf: &PyRef<Self>, _vm: &VirtualMachine) -> Cow<'static, PySequenceMethods> {
801+
static_cell! {
802+
static METHODS: PySequenceMethods;
803+
}
804+
Cow::Borrowed(METHODS.get_or_init(|| PySequenceMethods {
805+
length: Some(|zelf, _vm| Ok(zelf.payload::<Self>().unwrap().len())),
806+
concat: Some(|zelf, other, vm| {
807+
zelf.payload::<Self>()
808+
.unwrap()
809+
.inner()
810+
.concat(other, vm)
811+
.map(|x| PyByteArray::from(x).into_object(vm))
812+
}),
813+
repeat: Some(|zelf, n, vm| {
814+
zelf.payload::<Self>()
815+
.unwrap()
816+
.mul(n as isize, vm)
817+
.map(|x| x.into_object(vm))
818+
}),
819+
item: Some(|zelf, i, vm| zelf.payload::<Self>().unwrap().inner().item(i, vm)),
820+
ass_item: Some(|zelf, i, value, vm| {
821+
let zelf = zelf.payload::<Self>().unwrap();
822+
if let Some(value) = value {
823+
zelf.setitem_by_idx(i, value, vm)
824+
} else {
825+
zelf.delitem_by_idx(i, vm)
826+
}
827+
}),
828+
contains: Some(|zelf, other, vm| {
829+
let other = <Either<PyBytesInner, PyIntRef>>::try_from_object(vm, other.clone())?;
830+
zelf.payload::<Self>().unwrap().contains(other, vm)
831+
}),
832+
inplace_concat: Some(|zelf, other, vm| {
833+
let other = ArgBytesLike::try_from_object(vm, other.clone())?;
834+
let zelf = zelf.clone().downcast::<Self>().unwrap();
835+
Self::iadd(zelf, other, vm).map(|x| x.into())
836+
}),
837+
inplace_repeat: Some(|zelf, n, vm| {
838+
let zelf = zelf.clone().downcast::<Self>().unwrap();
839+
Self::imul(zelf, n as isize, vm).map(|x| x.into())
840+
}),
841+
}))
842+
}
843+
}
844+
793845
impl Unhashable for PyByteArray {}
794846

795847
impl Iterable for PyByteArray {

‎vm/src/builtins/bytes.rs

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -629,13 +629,7 @@ impl AsSequence for PyBytes {
629629
.new_bytes(zelf.payload::<Self>().unwrap().repeat(n))
630630
.into())
631631
}),
632-
item: Some(|zelf, i, vm| {
633-
zelf.payload::<Self>()
634-
.unwrap()
635-
.inner
636-
.item(i, vm)
637-
.map(|x| x.into_pyobject(vm))
638-
}),
632+
item: Some(|zelf, i, vm| zelf.payload::<Self>().unwrap().inner.item(i, vm)),
639633
contains: Some(|zelf, other, vm| {
640634
let other = <Either<PyBytesInner, PyIntRef>>::try_from_object(vm, other.clone())?;
641635
zelf.payload::<Self>().unwrap().contains(other, vm)

‎vm/src/bytesinner.rs

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -987,17 +987,12 @@ impl PyBytesInner {
987987
}
988988
}
989989

990-
pub fn item(&self, mut i: isize, vm: &VirtualMachine) -> PyResult<u8> {
991-
let len = self.len() as isize;
992-
if i < 0 {
993-
i += len;
994-
if i < 0 {
995-
return Err(vm.new_index_error("index out of range".to_string()));
996-
}
997-
} else if i >= len {
998-
return Err(vm.new_index_error("index out of range".to_string()));
999-
}
1000-
Ok(self.elements[i as usize])
990+
pub fn item(&self, i: isize, vm: &VirtualMachine) -> PyResult {
991+
let idx = self
992+
.elements
993+
.wrap_index(i)
994+
.ok_or_else(|| vm.new_index_error("index out of range".to_string()))?;
995+
Ok(vm.ctx.new_int(self.elements[idx]).into())
1001996
}
1002997
}
1003998

‎vm/src/protocol/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ mod buffer;
22
mod iter;
33
mod mapping;
44
mod object;
5-
pub(crate) mod sequence;
5+
mod sequence;
66

77
pub use buffer::{BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer};
88
pub use iter::{PyIter, PyIterIter, PyIterReturn};

‎vm/src/protocol/sequence.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,12 @@ pub struct PySequenceMethods {
1818
pub length: Option<fn(&PyObjectRef, &VirtualMachine) -> PyResult<usize>>,
1919
pub concat: Option<fn(&PyObjectRef, &PyObjectRef, &VirtualMachine) -> PyResult>,
2020
pub repeat: Option<fn(&PyObjectRef, usize, &VirtualMachine) -> PyResult>,
21-
pub inplace_concat: Option<fn(&PyObjectRef, &PyObjectRef, &VirtualMachine) -> PyResult>,
22-
pub inplace_repeat: Option<fn(&PyObjectRef, usize, &VirtualMachine) -> PyResult>,
2321
pub item: Option<fn(&PyObjectRef, isize, &VirtualMachine) -> PyResult>,
2422
pub ass_item:
2523
Option<fn(&PyObjectRef, isize, Option<PyObjectRef>, &VirtualMachine) -> PyResult<()>>,
2624
pub contains: Option<fn(&PyObjectRef, &PyObjectRef, &VirtualMachine) -> PyResult<bool>>,
25+
pub inplace_concat: Option<fn(&PyObjectRef, &PyObjectRef, &VirtualMachine) -> PyResult>,
26+
pub inplace_repeat: Option<fn(&PyObjectRef, usize, &VirtualMachine) -> PyResult>,
2727
}
2828

2929
impl PySequenceMethods {

0 commit comments

Comments
 (0)