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