rustpython_vm/builtins/
bytearray.rs

1//! Implementation of the python bytearray object.
2use super::{
3    PositionIterInternal, PyBytes, PyBytesRef, PyDictRef, PyIntRef, PyStrRef, PyTuple, PyTupleRef,
4    PyType, PyTypeRef,
5};
6use crate::{
7    anystr::{self, AnyStr},
8    atomic_func,
9    byte::{bytes_from_object, value_from_object},
10    bytesinner::{
11        bytes_decode, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
12        ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner,
13    },
14    class::PyClassImpl,
15    common::{
16        atomic::{AtomicUsize, Ordering},
17        lock::{
18            PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,
19            PyRwLockReadGuard, PyRwLockWriteGuard,
20        },
21    },
22    convert::{ToPyObject, ToPyResult},
23    function::{
24        ArgBytesLike, ArgIterable, ArgSize, Either, OptionalArg, OptionalOption, PyComparisonValue,
25    },
26    protocol::{
27        BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
28        PyMappingMethods, PyNumberMethods, PySequenceMethods,
29    },
30    sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
31    types::{
32        AsBuffer, AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor,
33        DefaultConstructor, Initializer, IterNext, Iterable, PyComparisonOp, Representable,
34        SelfIter, Unconstructible,
35    },
36    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
37    VirtualMachine,
38};
39use bstr::ByteSlice;
40use std::mem::size_of;
41
42#[pyclass(module = false, name = "bytearray", unhashable = true)]
43#[derive(Debug, Default)]
44pub struct PyByteArray {
45    inner: PyRwLock<PyBytesInner>,
46    exports: AtomicUsize,
47}
48
49pub type PyByteArrayRef = PyRef<PyByteArray>;
50
51impl From<PyBytesInner> for PyByteArray {
52    fn from(inner: PyBytesInner) -> Self {
53        Self::from_inner(inner)
54    }
55}
56
57impl From<Vec<u8>> for PyByteArray {
58    fn from(elements: Vec<u8>) -> Self {
59        Self::from(PyBytesInner { elements })
60    }
61}
62
63impl PyPayload for PyByteArray {
64    fn class(ctx: &Context) -> &'static Py<PyType> {
65        ctx.types.bytearray_type
66    }
67}
68
69/// Fill bytearray class methods dictionary.
70pub(crate) fn init(context: &Context) {
71    PyByteArray::extend_class(context, context.types.bytearray_type);
72    PyByteArrayIterator::extend_class(context, context.types.bytearray_iterator_type);
73}
74
75impl PyByteArray {
76    pub fn new_ref(data: Vec<u8>, ctx: &Context) -> PyRef<Self> {
77        PyRef::new_ref(Self::from(data), ctx.types.bytearray_type.to_owned(), None)
78    }
79
80    fn from_inner(inner: PyBytesInner) -> Self {
81        PyByteArray {
82            inner: PyRwLock::new(inner),
83            exports: AtomicUsize::new(0),
84        }
85    }
86
87    pub fn borrow_buf(&self) -> PyMappedRwLockReadGuard<'_, [u8]> {
88        PyRwLockReadGuard::map(self.inner.read(), |inner| &*inner.elements)
89    }
90
91    pub fn borrow_buf_mut(&self) -> PyMappedRwLockWriteGuard<'_, Vec<u8>> {
92        PyRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements)
93    }
94
95    fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
96        self.inner().mul(value, vm).map(|x| x.into())
97    }
98
99    fn _setitem_by_index(&self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
100        let value = value_from_object(vm, &value)?;
101        self.borrow_buf_mut().setitem_by_index(vm, i, value)
102    }
103
104    fn _setitem(
105        zelf: &Py<Self>,
106        needle: &PyObject,
107        value: PyObjectRef,
108        vm: &VirtualMachine,
109    ) -> PyResult<()> {
110        match SequenceIndex::try_from_borrowed_object(vm, needle, "bytearray")? {
111            SequenceIndex::Int(i) => zelf._setitem_by_index(i, value, vm),
112            SequenceIndex::Slice(slice) => {
113                let items = if zelf.is(&value) {
114                    zelf.borrow_buf().to_vec()
115                } else {
116                    bytes_from_object(vm, &value)?
117                };
118                if let Some(mut w) = zelf.try_resizable_opt() {
119                    w.elements.setitem_by_slice(vm, slice, &items)
120                } else {
121                    zelf.borrow_buf_mut()
122                        .setitem_by_slice_no_resize(vm, slice, &items)
123                }
124            }
125        }
126    }
127
128    fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
129        match SequenceIndex::try_from_borrowed_object(vm, needle, "bytearray")? {
130            SequenceIndex::Int(i) => self
131                .borrow_buf()
132                .getitem_by_index(vm, i)
133                .map(|x| vm.ctx.new_int(x).into()),
134            SequenceIndex::Slice(slice) => self
135                .borrow_buf()
136                .getitem_by_slice(vm, slice)
137                .map(|x| Self::new_ref(x, &vm.ctx).into()),
138        }
139    }
140
141    pub fn _delitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
142        match SequenceIndex::try_from_borrowed_object(vm, needle, "bytearray")? {
143            SequenceIndex::Int(i) => self.try_resizable(vm)?.elements.delitem_by_index(vm, i),
144            SequenceIndex::Slice(slice) => {
145                // TODO: delete 0 elements don't need resizable
146                self.try_resizable(vm)?.elements.delitem_by_slice(vm, slice)
147            }
148        }
149    }
150
151    fn irepeat(zelf: &Py<Self>, n: isize, vm: &VirtualMachine) -> PyResult<()> {
152        if n == 1 {
153            return Ok(());
154        }
155        let mut w = match zelf.try_resizable(vm) {
156            Ok(w) => w,
157            Err(err) => {
158                return if zelf.borrow_buf().is_empty() {
159                    // We can multiple an empty vector by any integer
160                    Ok(())
161                } else {
162                    Err(err)
163                };
164            }
165        };
166
167        w.imul(n, vm)
168    }
169}
170
171#[pyclass(
172    flags(BASETYPE),
173    with(
174        Py,
175        PyRef,
176        Constructor,
177        Initializer,
178        Comparable,
179        AsBuffer,
180        AsMapping,
181        AsSequence,
182        AsNumber,
183        Iterable,
184        Representable
185    )
186)]
187impl PyByteArray {
188    #[cfg(debug_assertions)]
189    #[pygetset]
190    fn exports(&self) -> usize {
191        self.exports.load(Ordering::Relaxed)
192    }
193
194    #[inline]
195    fn inner(&self) -> PyRwLockReadGuard<'_, PyBytesInner> {
196        self.inner.read()
197    }
198    #[inline]
199    fn inner_mut(&self) -> PyRwLockWriteGuard<'_, PyBytesInner> {
200        self.inner.write()
201    }
202
203    #[pymethod(magic)]
204    fn alloc(&self) -> usize {
205        self.inner().capacity()
206    }
207
208    #[pymethod(magic)]
209    fn len(&self) -> usize {
210        self.borrow_buf().len()
211    }
212
213    #[pymethod(magic)]
214    fn sizeof(&self) -> usize {
215        size_of::<Self>() + self.borrow_buf().len() * size_of::<u8>()
216    }
217
218    #[pymethod(magic)]
219    fn add(&self, other: ArgBytesLike) -> Self {
220        self.inner().add(&other.borrow_buf()).into()
221    }
222
223    #[pymethod(magic)]
224    fn contains(
225        &self,
226        needle: Either<PyBytesInner, PyIntRef>,
227        vm: &VirtualMachine,
228    ) -> PyResult<bool> {
229        self.inner().contains(needle, vm)
230    }
231
232    #[pymethod(magic)]
233    fn iadd(zelf: PyRef<Self>, other: ArgBytesLike, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
234        zelf.try_resizable(vm)?
235            .elements
236            .extend(&*other.borrow_buf());
237        Ok(zelf)
238    }
239
240    #[pymethod(magic)]
241    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
242        self._getitem(&needle, vm)
243    }
244
245    #[pymethod(magic)]
246    pub fn delitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
247        self._delitem(&needle, vm)
248    }
249
250    #[pystaticmethod]
251    fn maketrans(from: PyBytesInner, to: PyBytesInner, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
252        PyBytesInner::maketrans(from, to, vm)
253    }
254
255    #[pymethod]
256    fn isalnum(&self) -> bool {
257        self.inner().isalnum()
258    }
259
260    #[pymethod]
261    fn isalpha(&self) -> bool {
262        self.inner().isalpha()
263    }
264
265    #[pymethod]
266    fn isascii(&self) -> bool {
267        self.inner().isascii()
268    }
269
270    #[pymethod]
271    fn isdigit(&self) -> bool {
272        self.inner().isdigit()
273    }
274
275    #[pymethod]
276    fn islower(&self) -> bool {
277        self.inner().islower()
278    }
279
280    #[pymethod]
281    fn isspace(&self) -> bool {
282        self.inner().isspace()
283    }
284
285    #[pymethod]
286    fn isupper(&self) -> bool {
287        self.inner().isupper()
288    }
289
290    #[pymethod]
291    fn istitle(&self) -> bool {
292        self.inner().istitle()
293    }
294
295    #[pymethod]
296    fn lower(&self) -> Self {
297        self.inner().lower().into()
298    }
299
300    #[pymethod]
301    fn upper(&self) -> Self {
302        self.inner().upper().into()
303    }
304
305    #[pymethod]
306    fn capitalize(&self) -> Self {
307        self.inner().capitalize().into()
308    }
309
310    #[pymethod]
311    fn swapcase(&self) -> Self {
312        self.inner().swapcase().into()
313    }
314
315    #[pymethod]
316    fn hex(
317        &self,
318        sep: OptionalArg<Either<PyStrRef, PyBytesRef>>,
319        bytes_per_sep: OptionalArg<isize>,
320        vm: &VirtualMachine,
321    ) -> PyResult<String> {
322        self.inner().hex(sep, bytes_per_sep, vm)
323    }
324
325    #[pyclassmethod]
326    fn fromhex(cls: PyTypeRef, string: PyStrRef, vm: &VirtualMachine) -> PyResult {
327        let bytes = PyBytesInner::fromhex(string.as_str(), vm)?;
328        let bytes = vm.ctx.new_bytes(bytes);
329        let args = vec![bytes.into()].into();
330        PyType::call(&cls, args, vm)
331    }
332
333    #[pymethod]
334    fn center(
335        &self,
336        options: ByteInnerPaddingOptions,
337        vm: &VirtualMachine,
338    ) -> PyResult<PyByteArray> {
339        Ok(self.inner().center(options, vm)?.into())
340    }
341
342    #[pymethod]
343    fn ljust(
344        &self,
345        options: ByteInnerPaddingOptions,
346        vm: &VirtualMachine,
347    ) -> PyResult<PyByteArray> {
348        Ok(self.inner().ljust(options, vm)?.into())
349    }
350
351    #[pymethod]
352    fn rjust(
353        &self,
354        options: ByteInnerPaddingOptions,
355        vm: &VirtualMachine,
356    ) -> PyResult<PyByteArray> {
357        Ok(self.inner().rjust(options, vm)?.into())
358    }
359
360    #[pymethod]
361    fn count(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
362        self.inner().count(options, vm)
363    }
364
365    #[pymethod]
366    fn join(&self, iter: ArgIterable<PyBytesInner>, vm: &VirtualMachine) -> PyResult<PyByteArray> {
367        Ok(self.inner().join(iter, vm)?.into())
368    }
369
370    #[pymethod]
371    fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
372        let borrowed = self.borrow_buf();
373        let (affix, substr) =
374            match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
375                Some(x) => x,
376                None => return Ok(false),
377            };
378        substr.py_startsendswith(
379            &affix,
380            "endswith",
381            "bytes",
382            |s, x: PyBytesInner| s.ends_with(x.as_bytes()),
383            vm,
384        )
385    }
386
387    #[pymethod]
388    fn startswith(
389        &self,
390        options: anystr::StartsEndsWithArgs,
391        vm: &VirtualMachine,
392    ) -> PyResult<bool> {
393        let borrowed = self.borrow_buf();
394        let (affix, substr) =
395            match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
396                Some(x) => x,
397                None => return Ok(false),
398            };
399        substr.py_startsendswith(
400            &affix,
401            "startswith",
402            "bytes",
403            |s, x: PyBytesInner| s.starts_with(x.as_bytes()),
404            vm,
405        )
406    }
407
408    #[pymethod]
409    fn find(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
410        let index = self.inner().find(options, |h, n| h.find(n), vm)?;
411        Ok(index.map_or(-1, |v| v as isize))
412    }
413
414    #[pymethod]
415    fn index(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
416        let index = self.inner().find(options, |h, n| h.find(n), vm)?;
417        index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
418    }
419
420    #[pymethod]
421    fn rfind(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
422        let index = self.inner().find(options, |h, n| h.rfind(n), vm)?;
423        Ok(index.map_or(-1, |v| v as isize))
424    }
425
426    #[pymethod]
427    fn rindex(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
428        let index = self.inner().find(options, |h, n| h.rfind(n), vm)?;
429        index.ok_or_else(|| vm.new_value_error("substring not found".to_owned()))
430    }
431
432    #[pymethod]
433    fn translate(
434        &self,
435        options: ByteInnerTranslateOptions,
436        vm: &VirtualMachine,
437    ) -> PyResult<PyByteArray> {
438        Ok(self.inner().translate(options, vm)?.into())
439    }
440
441    #[pymethod]
442    fn strip(&self, chars: OptionalOption<PyBytesInner>) -> Self {
443        self.inner().strip(chars).into()
444    }
445
446    #[pymethod]
447    fn removeprefix(&self, prefix: PyBytesInner) -> Self {
448        self.inner().removeprefix(prefix).into()
449    }
450
451    #[pymethod]
452    fn removesuffix(&self, suffix: PyBytesInner) -> Self {
453        self.inner().removesuffix(suffix).to_vec().into()
454    }
455
456    #[pymethod]
457    fn split(
458        &self,
459        options: ByteInnerSplitOptions,
460        vm: &VirtualMachine,
461    ) -> PyResult<Vec<PyObjectRef>> {
462        self.inner().split(
463            options,
464            |s, vm| Self::new_ref(s.to_vec(), &vm.ctx).into(),
465            vm,
466        )
467    }
468
469    #[pymethod]
470    fn rsplit(
471        &self,
472        options: ByteInnerSplitOptions,
473        vm: &VirtualMachine,
474    ) -> PyResult<Vec<PyObjectRef>> {
475        self.inner().rsplit(
476            options,
477            |s, vm| Self::new_ref(s.to_vec(), &vm.ctx).into(),
478            vm,
479        )
480    }
481
482    #[pymethod]
483    fn partition(&self, sep: PyBytesInner, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
484        // sep ALWAYS converted to  bytearray even it's bytes or memoryview
485        // so its ok to accept PyBytesInner
486        let value = self.inner();
487        let (front, has_mid, back) = value.partition(&sep, vm)?;
488        Ok(vm.new_tuple((
489            Self::new_ref(front.to_vec(), &vm.ctx),
490            Self::new_ref(if has_mid { sep.elements } else { Vec::new() }, &vm.ctx),
491            Self::new_ref(back.to_vec(), &vm.ctx),
492        )))
493    }
494
495    #[pymethod]
496    fn rpartition(&self, sep: PyBytesInner, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
497        let value = self.inner();
498        let (back, has_mid, front) = value.rpartition(&sep, vm)?;
499        Ok(vm.new_tuple((
500            Self::new_ref(front.to_vec(), &vm.ctx),
501            Self::new_ref(if has_mid { sep.elements } else { Vec::new() }, &vm.ctx),
502            Self::new_ref(back.to_vec(), &vm.ctx),
503        )))
504    }
505
506    #[pymethod]
507    fn expandtabs(&self, options: anystr::ExpandTabsArgs) -> Self {
508        self.inner().expandtabs(options).into()
509    }
510
511    #[pymethod]
512    fn splitlines(&self, options: anystr::SplitLinesArgs, vm: &VirtualMachine) -> Vec<PyObjectRef> {
513        self.inner()
514            .splitlines(options, |x| Self::new_ref(x.to_vec(), &vm.ctx).into())
515    }
516
517    #[pymethod]
518    fn zfill(&self, width: isize) -> Self {
519        self.inner().zfill(width).into()
520    }
521
522    #[pymethod]
523    fn replace(
524        &self,
525        old: PyBytesInner,
526        new: PyBytesInner,
527        count: OptionalArg<isize>,
528        vm: &VirtualMachine,
529    ) -> PyResult<PyByteArray> {
530        Ok(self.inner().replace(old, new, count, vm)?.into())
531    }
532
533    #[pymethod]
534    fn copy(&self) -> Self {
535        self.borrow_buf().to_vec().into()
536    }
537
538    #[pymethod]
539    fn title(&self) -> Self {
540        self.inner().title().into()
541    }
542
543    #[pymethod(name = "__rmul__")]
544    #[pymethod(magic)]
545    fn mul(&self, value: ArgSize, vm: &VirtualMachine) -> PyResult<Self> {
546        self.repeat(value.into(), vm)
547    }
548
549    #[pymethod(magic)]
550    fn imul(zelf: PyRef<Self>, value: ArgSize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
551        Self::irepeat(&zelf, value.into(), vm)?;
552        Ok(zelf)
553    }
554
555    #[pymethod(name = "__mod__")]
556    fn mod_(&self, values: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyByteArray> {
557        let formatted = self.inner().cformat(values, vm)?;
558        Ok(formatted.into())
559    }
560
561    #[pymethod(magic)]
562    fn rmod(&self, _values: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
563        vm.ctx.not_implemented()
564    }
565
566    #[pymethod]
567    fn reverse(&self) {
568        self.borrow_buf_mut().reverse();
569    }
570}
571
572#[pyclass]
573impl Py<PyByteArray> {
574    #[pymethod(magic)]
575    fn setitem(
576        &self,
577        needle: PyObjectRef,
578        value: PyObjectRef,
579        vm: &VirtualMachine,
580    ) -> PyResult<()> {
581        PyByteArray::_setitem(self, &needle, value, vm)
582    }
583
584    #[pymethod]
585    fn pop(&self, index: OptionalArg<isize>, vm: &VirtualMachine) -> PyResult<u8> {
586        let elements = &mut self.try_resizable(vm)?.elements;
587        let index = elements
588            .wrap_index(index.unwrap_or(-1))
589            .ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
590        Ok(elements.remove(index))
591    }
592
593    #[pymethod]
594    fn insert(&self, index: isize, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
595        let value = value_from_object(vm, &object)?;
596        let elements = &mut self.try_resizable(vm)?.elements;
597        let index = elements.saturate_index(index);
598        elements.insert(index, value);
599        Ok(())
600    }
601
602    #[pymethod]
603    fn append(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
604        let value = value_from_object(vm, &object)?;
605        self.try_resizable(vm)?.elements.push(value);
606        Ok(())
607    }
608
609    #[pymethod]
610    fn remove(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
611        let value = value_from_object(vm, &object)?;
612        let elements = &mut self.try_resizable(vm)?.elements;
613        let index = elements
614            .find_byte(value)
615            .ok_or_else(|| vm.new_value_error("value not found in bytearray".to_owned()))?;
616        elements.remove(index);
617        Ok(())
618    }
619
620    #[pymethod]
621    fn extend(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
622        if self.is(&object) {
623            PyByteArray::irepeat(self, 2, vm)
624        } else {
625            let items = bytes_from_object(vm, &object)?;
626            self.try_resizable(vm)?.elements.extend(items);
627            Ok(())
628        }
629    }
630
631    #[pymethod]
632    fn clear(&self, vm: &VirtualMachine) -> PyResult<()> {
633        self.try_resizable(vm)?.elements.clear();
634        Ok(())
635    }
636
637    #[pymethod(magic)]
638    fn reduce_ex(
639        &self,
640        _proto: usize,
641        vm: &VirtualMachine,
642    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
643        Self::reduce(self, vm)
644    }
645
646    #[pymethod(magic)]
647    fn reduce(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
648        let bytes = PyBytes::from(self.borrow_buf().to_vec()).to_pyobject(vm);
649        (
650            self.class().to_owned(),
651            PyTuple::new_ref(vec![bytes], &vm.ctx),
652            self.as_object().dict(),
653        )
654    }
655}
656
657#[pyclass]
658impl PyRef<PyByteArray> {
659    #[pymethod]
660    fn lstrip(
661        self,
662        chars: OptionalOption<PyBytesInner>,
663        vm: &VirtualMachine,
664    ) -> PyRef<PyByteArray> {
665        let inner = self.inner();
666        let stripped = inner.lstrip(chars);
667        let elements = &inner.elements;
668        if stripped == elements {
669            drop(inner);
670            self
671        } else {
672            vm.ctx.new_pyref(PyByteArray::from(stripped.to_vec()))
673        }
674    }
675
676    #[pymethod]
677    fn rstrip(
678        self,
679        chars: OptionalOption<PyBytesInner>,
680        vm: &VirtualMachine,
681    ) -> PyRef<PyByteArray> {
682        let inner = self.inner();
683        let stripped = inner.rstrip(chars);
684        let elements = &inner.elements;
685        if stripped == elements {
686            drop(inner);
687            self
688        } else {
689            vm.ctx.new_pyref(PyByteArray::from(stripped.to_vec()))
690        }
691    }
692
693    #[pymethod]
694    fn decode(self, args: DecodeArgs, vm: &VirtualMachine) -> PyResult<PyStrRef> {
695        bytes_decode(self.into(), args, vm)
696    }
697}
698
699impl DefaultConstructor for PyByteArray {}
700
701impl Initializer for PyByteArray {
702    type Args = ByteInnerNewOptions;
703
704    fn init(zelf: PyRef<Self>, options: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
705        // First unpack bytearray and *then* get a lock to set it.
706        let mut inner = options.get_bytearray_inner(vm)?;
707        std::mem::swap(&mut *zelf.inner_mut(), &mut inner);
708        Ok(())
709    }
710}
711
712impl Comparable for PyByteArray {
713    fn cmp(
714        zelf: &Py<Self>,
715        other: &PyObject,
716        op: PyComparisonOp,
717        vm: &VirtualMachine,
718    ) -> PyResult<PyComparisonValue> {
719        if let Some(res) = op.identical_optimization(zelf, other) {
720            return Ok(res.into());
721        }
722        Ok(zelf.inner().cmp(other, op, vm))
723    }
724}
725
726static BUFFER_METHODS: BufferMethods = BufferMethods {
727    obj_bytes: |buffer| buffer.obj_as::<PyByteArray>().borrow_buf().into(),
728    obj_bytes_mut: |buffer| {
729        PyMappedRwLockWriteGuard::map(buffer.obj_as::<PyByteArray>().borrow_buf_mut(), |x| {
730            x.as_mut_slice()
731        })
732        .into()
733    },
734    release: |buffer| {
735        buffer
736            .obj_as::<PyByteArray>()
737            .exports
738            .fetch_sub(1, Ordering::Release);
739    },
740    retain: |buffer| {
741        buffer
742            .obj_as::<PyByteArray>()
743            .exports
744            .fetch_add(1, Ordering::Release);
745    },
746};
747
748impl AsBuffer for PyByteArray {
749    fn as_buffer(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
750        Ok(PyBuffer::new(
751            zelf.to_owned().into(),
752            BufferDescriptor::simple(zelf.len(), false),
753            &BUFFER_METHODS,
754        ))
755    }
756}
757
758impl BufferResizeGuard for PyByteArray {
759    type Resizable<'a> = PyRwLockWriteGuard<'a, PyBytesInner>;
760
761    fn try_resizable_opt(&self) -> Option<Self::Resizable<'_>> {
762        let w = self.inner.write();
763        (self.exports.load(Ordering::SeqCst) == 0).then_some(w)
764    }
765}
766
767impl AsMapping for PyByteArray {
768    fn as_mapping() -> &'static PyMappingMethods {
769        static AS_MAPPING: PyMappingMethods = PyMappingMethods {
770            length: atomic_func!(|mapping, _vm| Ok(PyByteArray::mapping_downcast(mapping).len())),
771            subscript: atomic_func!(|mapping, needle, vm| {
772                PyByteArray::mapping_downcast(mapping).getitem(needle.to_owned(), vm)
773            }),
774            ass_subscript: atomic_func!(|mapping, needle, value, vm| {
775                let zelf = PyByteArray::mapping_downcast(mapping);
776                if let Some(value) = value {
777                    Py::setitem(zelf, needle.to_owned(), value, vm)
778                } else {
779                    zelf.delitem(needle.to_owned(), vm)
780                }
781            }),
782        };
783        &AS_MAPPING
784    }
785}
786
787impl AsSequence for PyByteArray {
788    fn as_sequence() -> &'static PySequenceMethods {
789        static AS_SEQUENCE: PySequenceMethods = PySequenceMethods {
790            length: atomic_func!(|seq, _vm| Ok(PyByteArray::sequence_downcast(seq).len())),
791            concat: atomic_func!(|seq, other, vm| {
792                PyByteArray::sequence_downcast(seq)
793                    .inner()
794                    .concat(other, vm)
795                    .map(|x| PyByteArray::from(x).into_pyobject(vm))
796            }),
797            repeat: atomic_func!(|seq, n, vm| {
798                PyByteArray::sequence_downcast(seq)
799                    .repeat(n, vm)
800                    .map(|x| x.into_pyobject(vm))
801            }),
802            item: atomic_func!(|seq, i, vm| {
803                PyByteArray::sequence_downcast(seq)
804                    .borrow_buf()
805                    .getitem_by_index(vm, i)
806                    .map(|x| vm.ctx.new_bytes(vec![x]).into())
807            }),
808            ass_item: atomic_func!(|seq, i, value, vm| {
809                let zelf = PyByteArray::sequence_downcast(seq);
810                if let Some(value) = value {
811                    zelf._setitem_by_index(i, value, vm)
812                } else {
813                    zelf.borrow_buf_mut().delitem_by_index(vm, i)
814                }
815            }),
816            contains: atomic_func!(|seq, other, vm| {
817                let other =
818                    <Either<PyBytesInner, PyIntRef>>::try_from_object(vm, other.to_owned())?;
819                PyByteArray::sequence_downcast(seq).contains(other, vm)
820            }),
821            inplace_concat: atomic_func!(|seq, other, vm| {
822                let other = ArgBytesLike::try_from_object(vm, other.to_owned())?;
823                let zelf = PyByteArray::sequence_downcast(seq).to_owned();
824                PyByteArray::iadd(zelf, other, vm).map(|x| x.into())
825            }),
826            inplace_repeat: atomic_func!(|seq, n, vm| {
827                let zelf = PyByteArray::sequence_downcast(seq).to_owned();
828                PyByteArray::irepeat(&zelf, n, vm)?;
829                Ok(zelf.into())
830            }),
831        };
832        &AS_SEQUENCE
833    }
834}
835
836impl AsNumber for PyByteArray {
837    fn as_number() -> &'static PyNumberMethods {
838        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
839            remainder: Some(|a, b, vm| {
840                if let Some(a) = a.downcast_ref::<PyByteArray>() {
841                    a.mod_(b.to_owned(), vm).to_pyresult(vm)
842                } else {
843                    Ok(vm.ctx.not_implemented())
844                }
845            }),
846            ..PyNumberMethods::NOT_IMPLEMENTED
847        };
848        &AS_NUMBER
849    }
850}
851
852impl Iterable for PyByteArray {
853    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
854        Ok(PyByteArrayIterator {
855            internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
856        }
857        .into_pyobject(vm))
858    }
859}
860
861impl Representable for PyByteArray {
862    #[inline]
863    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
864        let class = zelf.class();
865        let class_name = class.name();
866        zelf.inner().repr_with_name(&class_name, vm)
867    }
868}
869
870// fn set_value(obj: &PyObject, value: Vec<u8>) {
871//     obj.borrow_mut().kind = PyObjectPayload::Bytes { value };
872// }
873
874#[pyclass(module = false, name = "bytearray_iterator")]
875#[derive(Debug)]
876pub struct PyByteArrayIterator {
877    internal: PyMutex<PositionIterInternal<PyByteArrayRef>>,
878}
879
880impl PyPayload for PyByteArrayIterator {
881    fn class(ctx: &Context) -> &'static Py<PyType> {
882        ctx.types.bytearray_iterator_type
883    }
884}
885
886#[pyclass(with(Unconstructible, IterNext, Iterable))]
887impl PyByteArrayIterator {
888    #[pymethod(magic)]
889    fn length_hint(&self) -> usize {
890        self.internal.lock().length_hint(|obj| obj.len())
891    }
892    #[pymethod(magic)]
893    fn reduce(&self, vm: &VirtualMachine) -> PyTupleRef {
894        self.internal
895            .lock()
896            .builtins_iter_reduce(|x| x.clone().into(), vm)
897    }
898
899    #[pymethod(magic)]
900    fn setstate(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
901        self.internal
902            .lock()
903            .set_state(state, |obj, pos| pos.min(obj.len()), vm)
904    }
905}
906
907impl Unconstructible for PyByteArrayIterator {}
908
909impl SelfIter for PyByteArrayIterator {}
910impl IterNext for PyByteArrayIterator {
911    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
912        zelf.internal.lock().next(|bytearray, pos| {
913            let buf = bytearray.borrow_buf();
914            Ok(PyIterReturn::from_result(
915                buf.get(pos).map(|&x| vm.new_pyobj(x)).ok_or(None),
916            ))
917        })
918    }
919}
Morty Proxy This is a proxified and sanitized view of the page, visit original site.