Skip to content

Commit 1db415d

Browse files
youknowoneyouknowone
committed
clean up bytes
1 parent 914e14b commit 1db415d

File tree

1 file changed

+78
-94
lines changed

1 file changed

+78
-94
lines changed

‎vm/src/builtins/bytes.rs‎

Lines changed: 78 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -101,23 +101,38 @@ impl PyBytes {
101101
PyRef::new_ref(Self::from(data), ctx.types.bytes_type.to_owned(), None)
102102
}
103103

104-
fn repeat(zelf: PyRef<Self>, count: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
105-
if count == 1 && zelf.class().is(vm.ctx.types.bytes_type) {
104+
fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
105+
match SequenceIndex::try_from_borrowed_object(vm, needle, "byte")? {
106+
SequenceIndex::Int(i) => self
107+
.getitem_by_index(vm, i)
108+
.map(|x| vm.ctx.new_int(x).into()),
109+
SequenceIndex::Slice(slice) => self
110+
.getitem_by_slice(vm, slice)
111+
.map(|x| vm.ctx.new_bytes(x).into()),
112+
}
113+
}
114+
}
115+
116+
impl PyRef<PyBytes> {
117+
fn repeat(self, count: isize, vm: &VirtualMachine) -> PyResult<PyRef<PyBytes>> {
118+
if count == 1 && self.class().is(vm.ctx.types.bytes_type) {
106119
// Special case: when some `bytes` is multiplied by `1`,
107120
// nothing really happens, we need to return an object itself
108121
// with the same `id()` to be compatible with CPython.
109122
// This only works for `bytes` itself, not its subclasses.
110-
return Ok(zelf);
123+
return Ok(self);
111124
}
112-
zelf.inner
125+
self.inner
113126
.mul(count, vm)
114-
.map(|x| Self::from(x).into_ref(vm))
127+
.map(|x| PyBytes::from(x).into_ref(vm))
115128
}
116129
}
117130

118131
#[pyclass(
119132
flags(BASETYPE),
120133
with(
134+
Py,
135+
PyRef,
121136
AsMapping,
122137
AsSequence,
123138
Hashable,
@@ -146,15 +161,6 @@ impl PyBytes {
146161
self.inner.as_bytes()
147162
}
148163

149-
#[pymethod(magic)]
150-
fn bytes(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
151-
if zelf.is(vm.ctx.types.bytes_type) {
152-
zelf
153-
} else {
154-
PyBytes::from(zelf.inner.clone()).into_ref(vm)
155-
}
156-
}
157-
158164
#[pymethod(magic)]
159165
fn sizeof(&self) -> usize {
160166
size_of::<Self>() + self.len() * size_of::<u8>()
@@ -179,17 +185,6 @@ impl PyBytes {
179185
PyBytesInner::maketrans(from, to, vm)
180186
}
181187

182-
fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
183-
match SequenceIndex::try_from_borrowed_object(vm, needle, "byte")? {
184-
SequenceIndex::Int(i) => self
185-
.getitem_by_index(vm, i)
186-
.map(|x| vm.ctx.new_int(x).into()),
187-
SequenceIndex::Slice(slice) => self
188-
.getitem_by_slice(vm, slice)
189-
.map(|x| vm.ctx.new_bytes(x).into()),
190-
}
191-
}
192-
193188
#[pymethod(magic)]
194189
fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
195190
self._getitem(&needle, vm)
@@ -371,53 +366,11 @@ impl PyBytes {
371366
self.inner.strip(chars).into()
372367
}
373368

374-
#[pymethod]
375-
fn lstrip(
376-
zelf: PyRef<Self>,
377-
chars: OptionalOption<PyBytesInner>,
378-
vm: &VirtualMachine,
379-
) -> PyRef<Self> {
380-
let stripped = zelf.inner.lstrip(chars);
381-
if stripped == zelf.as_bytes() {
382-
zelf
383-
} else {
384-
vm.ctx.new_bytes(stripped.to_vec())
385-
}
386-
}
387-
388-
#[pymethod]
389-
fn rstrip(
390-
zelf: PyRef<Self>,
391-
chars: OptionalOption<PyBytesInner>,
392-
vm: &VirtualMachine,
393-
) -> PyRef<Self> {
394-
let stripped = zelf.inner.rstrip(chars);
395-
if stripped == zelf.as_bytes() {
396-
zelf
397-
} else {
398-
vm.ctx.new_bytes(stripped.to_vec())
399-
}
400-
}
401-
402-
/// removeprefix($self, prefix, /)
403-
///
404-
///
405-
/// Return a bytes object with the given prefix string removed if present.
406-
///
407-
/// If the bytes starts with the prefix string, return string[len(prefix):]
408-
/// Otherwise, return a copy of the original bytes.
409369
#[pymethod]
410370
fn removeprefix(&self, prefix: PyBytesInner) -> Self {
411371
self.inner.removeprefix(prefix).into()
412372
}
413373

414-
/// removesuffix(self, prefix, /)
415-
///
416-
///
417-
/// Return a bytes object with the given suffix string removed if present.
418-
///
419-
/// If the bytes ends with the suffix string, return string[:len(suffix)]
420-
/// Otherwise, return a copy of the original bytes.
421374
#[pymethod]
422375
fn removesuffix(&self, suffix: PyBytesInner) -> Self {
423376
self.inner.removesuffix(suffix).into()
@@ -508,7 +461,7 @@ impl PyBytes {
508461
#[pymethod(name = "__rmul__")]
509462
#[pymethod(magic)]
510463
fn mul(zelf: PyRef<Self>, value: ArgIndex, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
511-
Self::repeat(zelf, value.try_to_primitive(vm)?, vm)
464+
zelf.repeat(value.try_to_primitive(vm)?, vm)
512465
}
513466

514467
#[pymethod(name = "__mod__")]
@@ -522,47 +475,79 @@ impl PyBytes {
522475
vm.ctx.not_implemented()
523476
}
524477

525-
/// Return a string decoded from the given bytes.
526-
/// Default encoding is 'utf-8'.
527-
/// Default errors is 'strict', meaning that encoding errors raise a UnicodeError.
528-
/// Other possible values are 'ignore', 'replace'
529-
/// For a list of possible encodings,
530-
/// see https://docs.python.org/3/library/codecs.html#standard-encodings
531-
/// currently, only 'utf-8' and 'ascii' emplemented
532-
#[pymethod]
533-
fn decode(zelf: PyRef<Self>, args: DecodeArgs, vm: &VirtualMachine) -> PyResult<PyStrRef> {
534-
bytes_decode(zelf.into(), args, vm)
535-
}
536-
537478
#[pymethod(magic)]
538479
fn getnewargs(&self, vm: &VirtualMachine) -> PyTupleRef {
539480
let param: Vec<PyObjectRef> = self.elements().map(|x| x.to_pyobject(vm)).collect();
540481
PyTuple::new_ref(param, &vm.ctx)
541482
}
483+
}
542484

485+
#[pyclass]
486+
impl Py<PyBytes> {
543487
#[pymethod(magic)]
544488
fn reduce_ex(
545-
zelf: PyRef<Self>,
489+
&self,
546490
_proto: usize,
547491
vm: &VirtualMachine,
548492
) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
549-
Self::reduce(zelf, vm)
493+
Self::reduce(self, vm)
550494
}
551495

552496
#[pymethod(magic)]
553-
fn reduce(
554-
zelf: PyRef<Self>,
555-
vm: &VirtualMachine,
556-
) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
557-
let bytes = PyBytes::from(zelf.to_vec()).to_pyobject(vm);
497+
fn reduce(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
498+
let bytes = PyBytes::from(self.to_vec()).to_pyobject(vm);
558499
(
559-
zelf.class().to_owned(),
500+
self.class().to_owned(),
560501
PyTuple::new_ref(vec![bytes], &vm.ctx),
561-
zelf.as_object().dict(),
502+
self.as_object().dict(),
562503
)
563504
}
564505
}
565506

507+
#[pyclass]
508+
impl PyRef<PyBytes> {
509+
#[pymethod(magic)]
510+
fn bytes(self, vm: &VirtualMachine) -> PyRef<PyBytes> {
511+
if self.is(vm.ctx.types.bytes_type) {
512+
self
513+
} else {
514+
PyBytes::from(self.inner.clone()).into_ref(vm)
515+
}
516+
}
517+
518+
#[pymethod]
519+
fn lstrip(self, chars: OptionalOption<PyBytesInner>, vm: &VirtualMachine) -> PyRef<PyBytes> {
520+
let stripped = self.inner.lstrip(chars);
521+
if stripped == self.as_bytes() {
522+
self
523+
} else {
524+
vm.ctx.new_bytes(stripped.to_vec())
525+
}
526+
}
527+
528+
#[pymethod]
529+
fn rstrip(self, chars: OptionalOption<PyBytesInner>, vm: &VirtualMachine) -> PyRef<PyBytes> {
530+
let stripped = self.inner.rstrip(chars);
531+
if stripped == self.as_bytes() {
532+
self
533+
} else {
534+
vm.ctx.new_bytes(stripped.to_vec())
535+
}
536+
}
537+
538+
/// Return a string decoded from the given bytes.
539+
/// Default encoding is 'utf-8'.
540+
/// Default errors is 'strict', meaning that encoding errors raise a UnicodeError.
541+
/// Other possible values are 'ignore', 'replace'
542+
/// For a list of possible encodings,
543+
/// see https://docs.python.org/3/library/codecs.html#standard-encodings
544+
/// currently, only 'utf-8' and 'ascii' emplemented
545+
#[pymethod]
546+
fn decode(self, args: DecodeArgs, vm: &VirtualMachine) -> PyResult<PyStrRef> {
547+
bytes_decode(self.into(), args, vm)
548+
}
549+
}
550+
566551
static BUFFER_METHODS: BufferMethods = BufferMethods {
567552
obj_bytes: |buffer| buffer.obj_as::<PyBytes>().as_bytes().into(),
568553
obj_bytes_mut: |_| panic!(),
@@ -605,11 +590,10 @@ impl AsSequence for PyBytes {
605590
.map(|x| vm.ctx.new_bytes(x).into())
606591
}),
607592
repeat: atomic_func!(|seq, n, vm| {
608-
if let Ok(zelf) = seq.obj.to_owned().downcast::<PyBytes>() {
609-
PyBytes::repeat(zelf, n, vm).to_pyresult(vm)
610-
} else {
611-
Err(vm.new_type_error("bad argument type for built-in operation".to_owned()))
612-
}
593+
let zelf = seq.obj.to_owned().downcast::<PyBytes>().map_err(|_| {
594+
vm.new_type_error("bad argument type for built-in operation".to_owned())
595+
})?;
596+
zelf.repeat(n, vm).to_pyresult(vm)
613597
}),
614598
item: atomic_func!(|seq, i, vm| {
615599
PyBytes::sequence_downcast(seq)

0 commit comments

Comments
 (0)