Skip to content

Commit aa1591f

Browse files
youknowoneyouknowone
committed
bytes
1 parent f9e2fb4 commit aa1591f

File tree

3 files changed

+64
-56
lines changed

3 files changed

+64
-56
lines changed

‎crates/vm/src/builtins/bytes.rs‎

Lines changed: 56 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -92,15 +92,66 @@ pub(crate) fn init(context: &Context) {
9292
}
9393

9494
impl Constructor for PyBytes {
95-
type Args = ByteInnerNewOptions;
95+
type Args = Vec<u8>;
9696

9797
fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
98-
let options: Self::Args = args.bind(vm)?;
99-
options.get_bytes(cls, vm).to_pyresult(vm)
98+
let options: ByteInnerNewOptions = args.bind(vm)?;
99+
100+
// Optimizations for exact bytes type
101+
if cls.is(vm.ctx.types.bytes_type) {
102+
// Return empty bytes singleton
103+
if options.source.is_missing()
104+
&& options.encoding.is_missing()
105+
&& options.errors.is_missing()
106+
{
107+
return Ok(vm.ctx.empty_bytes.clone().into());
108+
}
109+
110+
// Return exact bytes as-is
111+
if let OptionalArg::Present(ref obj) = options.source
112+
&& options.encoding.is_missing()
113+
&& options.errors.is_missing()
114+
{
115+
if let Ok(b) = obj.clone().downcast_exact::<PyBytes>(vm) {
116+
return Ok(b.into_pyref().into());
117+
}
118+
}
119+
}
120+
121+
// Handle __bytes__ method - may return PyBytes directly
122+
if let OptionalArg::Present(ref obj) = options.source
123+
&& options.encoding.is_missing()
124+
&& options.errors.is_missing()
125+
{
126+
if let Some(bytes_method) = vm.get_method(obj.clone(), identifier!(vm, __bytes__)) {
127+
let bytes = bytes_method?.call((), vm)?;
128+
// If exact bytes type and __bytes__ returns bytes, use it directly
129+
if cls.is(vm.ctx.types.bytes_type) {
130+
if let Ok(b) = bytes.clone().downcast::<PyBytes>() {
131+
return Ok(b.into());
132+
}
133+
}
134+
// Otherwise convert to Vec<u8>
135+
let inner = PyBytesInner::try_from_borrowed_object(vm, &bytes)?;
136+
let payload = Self::py_new(&cls, inner.elements, vm)?;
137+
return payload.into_ref_with_type(vm, cls).map(Into::into);
138+
}
139+
}
140+
141+
// Fallback to get_bytearray_inner
142+
let elements = options.get_bytearray_inner(vm)?.elements;
143+
144+
// Return empty bytes singleton for exact bytes types
145+
if elements.is_empty() && cls.is(vm.ctx.types.bytes_type) {
146+
return Ok(vm.ctx.empty_bytes.clone().into());
147+
}
148+
149+
let payload = Self::py_new(&cls, elements, vm)?;
150+
payload.into_ref_with_type(vm, cls).map(Into::into)
100151
}
101152

102-
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
103-
unreachable!("use slot_new")
153+
fn py_new(_cls: &Py<PyType>, elements: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
154+
Ok(Self::from(elements))
104155
}
105156
}
106157

‎crates/vm/src/bytes_inner.rs‎

Lines changed: 2 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
11
// spell-checker:ignore unchunked
22
use crate::{
3-
AsObject, PyObject, PyObjectRef, PyPayload, PyResult, TryFromBorrowedObject, VirtualMachine,
3+
AsObject, PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, VirtualMachine,
44
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
55
builtins::{
66
PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyStr, PyStrRef,
7-
PyTypeRef, pystr,
7+
pystr,
88
},
99
byte::bytes_from_object,
1010
cformat::cformat_bytes,
1111
common::hash,
1212
function::{ArgIterable, Either, OptionalArg, OptionalOption, PyComparisonValue},
13-
identifier,
1413
literal::escape::Escape,
1514
protocol::PyBuffer,
1615
sequence::{SequenceExt, SequenceMutExt},
@@ -91,43 +90,6 @@ impl ByteInnerNewOptions {
9190
})
9291
}
9392

94-
pub fn get_bytes(self, cls: PyTypeRef, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
95-
let inner = match (&self.source, &self.encoding, &self.errors) {
96-
(OptionalArg::Present(obj), OptionalArg::Missing, OptionalArg::Missing) => {
97-
let obj = obj.clone();
98-
// construct an exact bytes from an exact bytes do not clone
99-
let obj = if cls.is(vm.ctx.types.bytes_type) {
100-
match obj.downcast_exact::<PyBytes>(vm) {
101-
Ok(b) => return Ok(b.into_pyref()),
102-
Err(obj) => obj,
103-
}
104-
} else {
105-
obj
106-
};
107-
108-
if let Some(bytes_method) = vm.get_method(obj, identifier!(vm, __bytes__)) {
109-
// construct an exact bytes from __bytes__ slot.
110-
// if __bytes__ return a bytes, use the bytes object except we are the subclass of the bytes
111-
let bytes = bytes_method?.call((), vm)?;
112-
let bytes = if cls.is(vm.ctx.types.bytes_type) {
113-
match bytes.downcast::<PyBytes>() {
114-
Ok(b) => return Ok(b),
115-
Err(bytes) => bytes,
116-
}
117-
} else {
118-
bytes
119-
};
120-
Some(PyBytesInner::try_from_borrowed_object(vm, &bytes))
121-
} else {
122-
None
123-
}
124-
}
125-
_ => None,
126-
}
127-
.unwrap_or_else(|| self.get_bytearray_inner(vm))?;
128-
PyBytes::from(inner).into_ref_with_type(vm, cls)
129-
}
130-
13193
pub fn get_bytearray_inner(self, vm: &VirtualMachine) -> PyResult<PyBytesInner> {
13294
match (self.source, self.encoding, self.errors) {
13395
(OptionalArg::Present(obj), OptionalArg::Missing, OptionalArg::Missing) => {

‎crates/vm/src/protocol/object.rs‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,16 @@
44
use crate::{
55
AsObject, Py, PyObject, PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine,
66
builtins::{
7-
PyAsyncGen, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, PyStr, PyTuple, PyTupleRef,
8-
PyType, PyTypeRef, PyUtf8Str, pystr::AsPyStr,
7+
PyAsyncGen, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, PyStr, PyTuple,
8+
PyTupleRef, PyType, PyTypeRef, PyUtf8Str, pystr::AsPyStr,
99
},
10-
bytes_inner::ByteInnerNewOptions,
1110
common::{hash::PyHash, str::to_ascii},
1211
convert::{ToPyObject, ToPyResult},
1312
dict_inner::DictKey,
14-
function::{Either, OptionalArg, PyArithmeticValue, PySetterValue},
13+
function::{Either, FuncArgs, PyArithmeticValue, PySetterValue},
1514
object::PyPayload,
1615
protocol::{PyIter, PyMapping, PySequence},
17-
types::PyComparisonOp,
16+
types::{Constructor, PyComparisonOp},
1817
};
1918

2019
// RustPython doesn't need these items
@@ -37,12 +36,8 @@ impl PyObjectRef {
3736
match self.downcast_exact::<PyInt>(vm) {
3837
Ok(int) => Err(vm.new_downcast_type_error(bytes_type, &int)),
3938
Err(obj) => {
40-
let options = ByteInnerNewOptions {
41-
source: OptionalArg::Present(obj),
42-
encoding: OptionalArg::Missing,
43-
errors: OptionalArg::Missing,
44-
};
45-
options.get_bytes(bytes_type.to_owned(), vm).map(Into::into)
39+
let args = FuncArgs::from(vec![obj]);
40+
<PyBytes as Constructor>::slot_new(bytes_type.to_owned(), args, vm)
4641
}
4742
}
4843
}

0 commit comments

Comments
 (0)