Skip to main content

rustpython_vm/builtins/
bytearray.rs

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