1use 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
69pub(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 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 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 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 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 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}