|
1 | 1 | use crate::common::borrow::{BorrowedValue, BorrowedValueMut}; |
2 | | -use crate::common::lock::{PyRwLockReadGuard, PyRwLockWriteGuard}; |
3 | | -use crate::obj::objbytearray::{PyByteArray, PyByteArrayRef}; |
4 | | -use crate::obj::objbytes::{PyBytes, PyBytesRef}; |
5 | | -use crate::pyobject::PyObjectRef; |
6 | | -use crate::pyobject::{BorrowValue, PyResult, TryFromObject, TypeProtocol}; |
7 | | -use crate::stdlib::array::{PyArray, PyArrayRef}; |
| 2 | +use crate::obj::objmemory::{try_buffer_from_object, BufferRef}; |
| 3 | +use crate::pyobject::{BorrowValue, PyObjectRef, PyResult, TryFromObject}; |
8 | 4 | use crate::vm::VirtualMachine; |
9 | 5 |
|
10 | 6 | #[derive(Debug)] |
11 | | -pub enum PyBytesLike { |
12 | | - Bytes(PyBytesRef), |
13 | | - Bytearray(PyByteArrayRef), |
14 | | - Array(PyArrayRef), |
15 | | -} |
| 7 | +pub struct PyBytesLike(BufferRef); |
16 | 8 |
|
17 | | -impl TryFromObject for PyBytesLike { |
18 | | - fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> { |
19 | | - match_class!(match obj { |
20 | | - b @ PyBytes => Ok(PyBytesLike::Bytes(b)), |
21 | | - b @ PyByteArray => Ok(PyBytesLike::Bytearray(b)), |
22 | | - array @ PyArray => Ok(PyBytesLike::Array(array)), |
23 | | - obj => Err(vm.new_type_error(format!( |
24 | | - "a bytes-like object is required, not {}", |
25 | | - obj.lease_class().name |
26 | | - ))), |
27 | | - }) |
| 9 | +#[derive(Debug)] |
| 10 | +pub struct PyRwBytesLike(BufferRef); |
| 11 | + |
| 12 | +impl PyBytesLike { |
| 13 | + pub fn with_ref<F, R>(&self, f: F) -> R |
| 14 | + where |
| 15 | + F: FnOnce(&[u8]) -> R, |
| 16 | + { |
| 17 | + f(&*self.borrow_value()) |
28 | 18 | } |
29 | | -} |
30 | 19 |
|
31 | | -impl<'a> BorrowValue<'a> for PyBytesLike { |
32 | | - type Borrowed = BorrowedValue<'a, [u8]>; |
33 | | - fn borrow_value(&'a self) -> Self::Borrowed { |
34 | | - match self { |
35 | | - Self::Bytes(b) => b.borrow_value().into(), |
36 | | - Self::Bytearray(b) => { |
37 | | - PyRwLockReadGuard::map(b.borrow_value(), |b| b.elements.as_slice()).into() |
38 | | - } |
39 | | - Self::Array(a) => a.get_bytes().into(), |
40 | | - } |
| 20 | + pub fn len(&self) -> usize { |
| 21 | + self.borrow_value().len() |
| 22 | + } |
| 23 | + |
| 24 | + pub fn is_empty(&self) -> bool { |
| 25 | + self.borrow_value().is_empty() |
| 26 | + } |
| 27 | + |
| 28 | + pub fn to_cow(&self) -> std::borrow::Cow<[u8]> { |
| 29 | + self.borrow_value().to_vec().into() |
41 | 30 | } |
42 | 31 | } |
43 | 32 |
|
44 | | -impl PyBytesLike { |
| 33 | +impl PyRwBytesLike { |
| 34 | + pub fn with_ref<F, R>(&self, f: F) -> R |
| 35 | + where |
| 36 | + F: FnOnce(&mut [u8]) -> R, |
| 37 | + { |
| 38 | + f(&mut *self.borrow_value()) |
| 39 | + } |
| 40 | + |
45 | 41 | pub fn len(&self) -> usize { |
46 | | - match self { |
47 | | - PyBytesLike::Bytes(b) => b.len(), |
48 | | - PyBytesLike::Bytearray(b) => b.borrow_value().len(), |
49 | | - PyBytesLike::Array(array) => array.len(), |
50 | | - } |
| 42 | + self.borrow_value().len() |
51 | 43 | } |
52 | 44 |
|
53 | 45 | pub fn is_empty(&self) -> bool { |
54 | | - self.len() == 0 |
| 46 | + self.borrow_value().is_empty() |
55 | 47 | } |
| 48 | +} |
56 | 49 |
|
57 | | - pub fn to_cow(&self) -> std::borrow::Cow<[u8]> { |
58 | | - match self { |
59 | | - PyBytesLike::Bytes(b) => b.borrow_value().into(), |
60 | | - PyBytesLike::Bytearray(b) => b.borrow_value().elements.clone().into(), |
61 | | - PyBytesLike::Array(array) => array.tobytes().into(), |
| 50 | +impl TryFromObject for PyBytesLike { |
| 51 | + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> { |
| 52 | + let buffer = try_buffer_from_object(vm, &obj)?; |
| 53 | + if buffer.get_options().contiguous { |
| 54 | + Ok(Self(buffer)) |
| 55 | + } else { |
| 56 | + Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())) |
62 | 57 | } |
63 | 58 | } |
| 59 | +} |
64 | 60 |
|
65 | | - #[inline] |
66 | | - pub fn with_ref<R>(&self, f: impl FnOnce(&[u8]) -> R) -> R { |
67 | | - match self { |
68 | | - PyBytesLike::Bytes(b) => f(b.borrow_value()), |
69 | | - PyBytesLike::Bytearray(b) => f(&b.borrow_value().elements), |
70 | | - PyBytesLike::Array(array) => f(&*array.get_bytes()), |
71 | | - } |
| 61 | +impl<'a> BorrowValue<'a> for PyBytesLike { |
| 62 | + type Borrowed = BorrowedValue<'a, [u8]>; |
| 63 | + fn borrow_value(&'a self) -> Self::Borrowed { |
| 64 | + self.0.as_contiguous().unwrap() |
72 | 65 | } |
73 | 66 | } |
74 | 67 |
|
75 | | -pub(crate) fn try_bytes_like<R>( |
| 68 | +pub fn try_bytes_like<R>( |
76 | 69 | vm: &VirtualMachine, |
77 | 70 | obj: &PyObjectRef, |
78 | 71 | f: impl FnOnce(&[u8]) -> R, |
79 | 72 | ) -> PyResult<R> { |
80 | | - let r = match_class!(match obj { |
81 | | - ref b @ PyBytes => f(b.borrow_value()), |
82 | | - ref b @ PyByteArray => f(&b.borrow_value().elements), |
83 | | - ref array @ PyArray => f(&*array.get_bytes()), |
84 | | - obj => |
85 | | - return Err(vm.new_type_error(format!( |
86 | | - "a bytes-like object is required, not {}", |
87 | | - obj.lease_class().name |
88 | | - ))), |
89 | | - }); |
90 | | - Ok(r) |
| 73 | + let buffer = try_buffer_from_object(vm, obj)?; |
| 74 | + buffer.as_contiguous().map(|x| f(&*x)).ok_or_else(|| { |
| 75 | + vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned()) |
| 76 | + }) |
91 | 77 | } |
92 | 78 |
|
93 | | -pub enum PyRwBytesLike { |
94 | | - Bytearray(PyByteArrayRef), |
95 | | - Array(PyArrayRef), |
| 79 | +pub fn try_rw_bytes_like<R>( |
| 80 | + vm: &VirtualMachine, |
| 81 | + obj: &PyObjectRef, |
| 82 | + f: impl FnOnce(&mut [u8]) -> R, |
| 83 | +) -> PyResult<R> { |
| 84 | + let buffer = try_buffer_from_object(vm, obj)?; |
| 85 | + buffer |
| 86 | + .as_contiguous_mut() |
| 87 | + .map(|mut x| f(&mut *x)) |
| 88 | + .ok_or_else(|| vm.new_type_error("buffer is not a read-write bytes-like object".to_owned())) |
96 | 89 | } |
97 | 90 |
|
98 | 91 | impl TryFromObject for PyRwBytesLike { |
99 | 92 | fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> { |
100 | | - match_class!(match obj { |
101 | | - b @ PyByteArray => Ok(PyRwBytesLike::Bytearray(b)), |
102 | | - array @ PyArray => Ok(PyRwBytesLike::Array(array)), |
103 | | - obj => Err(vm.new_type_error(format!( |
104 | | - "a buffer object is required, not '{}'", |
105 | | - obj.lease_class().name |
106 | | - ))), |
107 | | - }) |
| 93 | + let buffer = try_buffer_from_object(vm, &obj)?; |
| 94 | + if !buffer.get_options().contiguous { |
| 95 | + Err(vm.new_type_error("non-contiguous buffer is not a bytes-like object".to_owned())) |
| 96 | + } else if buffer.get_options().readonly { |
| 97 | + Err(vm.new_type_error("buffer is not a read-write bytes-like object".to_owned())) |
| 98 | + } else { |
| 99 | + Ok(Self(buffer)) |
| 100 | + } |
108 | 101 | } |
109 | 102 | } |
110 | 103 |
|
111 | 104 | impl<'a> BorrowValue<'a> for PyRwBytesLike { |
112 | 105 | type Borrowed = BorrowedValueMut<'a, [u8]>; |
113 | 106 | fn borrow_value(&'a self) -> Self::Borrowed { |
114 | | - match self { |
115 | | - Self::Bytearray(b) => { |
116 | | - PyRwLockWriteGuard::map(b.borrow_value_mut(), |b| b.elements.as_mut_slice()).into() |
117 | | - } |
118 | | - Self::Array(a) => a.get_bytes_mut().into(), |
119 | | - } |
120 | | - } |
121 | | -} |
122 | | - |
123 | | -impl PyRwBytesLike { |
124 | | - pub fn len(&self) -> usize { |
125 | | - match self { |
126 | | - PyRwBytesLike::Bytearray(b) => b.borrow_value().len(), |
127 | | - PyRwBytesLike::Array(array) => array.len(), |
128 | | - } |
129 | | - } |
130 | | - |
131 | | - pub fn is_empty(&self) -> bool { |
132 | | - self.len() == 0 |
133 | | - } |
134 | | - |
135 | | - #[inline] |
136 | | - pub fn with_ref<R>(&self, f: impl FnOnce(&mut [u8]) -> R) -> R { |
137 | | - match self { |
138 | | - PyRwBytesLike::Bytearray(b) => f(&mut b.borrow_value_mut().elements), |
139 | | - PyRwBytesLike::Array(array) => f(&mut array.get_bytes_mut()), |
140 | | - } |
| 107 | + self.0.as_contiguous_mut().unwrap() |
141 | 108 | } |
142 | 109 | } |
0 commit comments