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