Skip to content

Commit 9af05c3

Browse files
committed
AsBuffer
1 parent 8556955 commit 9af05c3

File tree

7 files changed

+679
-300
lines changed

7 files changed

+679
-300
lines changed

crates/vm/src/stdlib/ctypes.rs

Lines changed: 161 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
4545

4646
#[pymodule]
4747
pub(crate) mod _ctypes {
48-
use super::base::PyCSimple;
48+
use super::base::{CDataObject, PyCSimple};
4949
use crate::builtins::PyTypeRef;
5050
use crate::class::StaticType;
5151
use crate::convert::ToPyObject;
@@ -55,7 +55,7 @@ pub(crate) mod _ctypes {
5555
use crossbeam_utils::atomic::AtomicCell;
5656
use std::ffi::{
5757
c_double, c_float, c_int, c_long, c_longlong, c_schar, c_short, c_uchar, c_uint, c_ulong,
58-
c_ulonglong,
58+
c_ulonglong, c_ushort,
5959
};
6060
use std::mem;
6161
use widestring::WideChar;
@@ -159,14 +159,15 @@ pub(crate) mod _ctypes {
159159
}
160160

161161
/// Get the size of a ctypes type from its type object
162+
#[allow(dead_code)]
162163
pub fn get_size_from_type(cls: &PyTypeRef, vm: &VirtualMachine) -> PyResult<usize> {
163164
// Try to get _type_ attribute for simple types
164-
if let Ok(type_attr) = cls.as_object().get_attr("_type_", vm) {
165-
if let Ok(s) = type_attr.str(vm) {
166-
let s = s.to_string();
167-
if s.len() == 1 && SIMPLE_TYPE_CHARS.contains(s.as_str()) {
168-
return Ok(get_size(&s));
169-
}
165+
if let Ok(type_attr) = cls.as_object().get_attr("_type_", vm)
166+
&& let Ok(s) = type_attr.str(vm)
167+
{
168+
let s = s.to_string();
169+
if s.len() == 1 && SIMPLE_TYPE_CHARS.contains(s.as_str()) {
170+
return Ok(get_size(&s));
170171
}
171172
}
172173
// Fall back to sizeof
@@ -180,138 +181,159 @@ pub(crate) mod _ctypes {
180181
vm: &VirtualMachine,
181182
) -> PyResult<PyObjectRef> {
182183
// Try to get _type_ attribute
183-
if let Ok(type_attr) = cls.as_object().get_attr("_type_", vm) {
184-
if let Ok(s) = type_attr.str(vm) {
185-
let ty = s.to_string();
186-
return match ty.as_str() {
187-
"c" => {
188-
// c_char - single byte
189-
Ok(vm.ctx.new_bytes(bytes.to_vec()).into())
190-
}
191-
"b" => {
192-
// c_byte - signed char
193-
let val = if !bytes.is_empty() { bytes[0] as i8 } else { 0 };
194-
Ok(vm.ctx.new_int(val).into())
195-
}
196-
"B" => {
197-
// c_ubyte - unsigned char
198-
let val = if !bytes.is_empty() { bytes[0] } else { 0 };
199-
Ok(vm.ctx.new_int(val).into())
200-
}
201-
"h" => {
202-
// c_short
203-
let val = if bytes.len() >= 2 {
204-
i16::from_ne_bytes([bytes[0], bytes[1]])
205-
} else {
206-
0
207-
};
208-
Ok(vm.ctx.new_int(val).into())
209-
}
210-
"H" => {
211-
// c_ushort
212-
let val = if bytes.len() >= 2 {
213-
u16::from_ne_bytes([bytes[0], bytes[1]])
214-
} else {
215-
0
216-
};
217-
Ok(vm.ctx.new_int(val).into())
218-
}
219-
"i" | "l" => {
220-
// c_int, c_long (assuming 32-bit)
221-
let val = if bytes.len() >= 4 {
222-
i32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
184+
if let Ok(type_attr) = cls.as_object().get_attr("_type_", vm)
185+
&& let Ok(s) = type_attr.str(vm)
186+
{
187+
let ty = s.to_string();
188+
return match ty.as_str() {
189+
"c" => {
190+
// c_char - single byte
191+
Ok(vm.ctx.new_bytes(bytes.to_vec()).into())
192+
}
193+
"b" => {
194+
// c_byte - signed char
195+
let val = if !bytes.is_empty() { bytes[0] as i8 } else { 0 };
196+
Ok(vm.ctx.new_int(val).into())
197+
}
198+
"B" => {
199+
// c_ubyte - unsigned char
200+
let val = if !bytes.is_empty() { bytes[0] } else { 0 };
201+
Ok(vm.ctx.new_int(val).into())
202+
}
203+
"h" => {
204+
// c_short
205+
const SIZE: usize = mem::size_of::<c_short>();
206+
let val = if bytes.len() >= SIZE {
207+
c_short::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
208+
} else {
209+
0
210+
};
211+
Ok(vm.ctx.new_int(val).into())
212+
}
213+
"H" => {
214+
// c_ushort
215+
const SIZE: usize = mem::size_of::<c_ushort>();
216+
let val = if bytes.len() >= SIZE {
217+
c_ushort::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
218+
} else {
219+
0
220+
};
221+
Ok(vm.ctx.new_int(val).into())
222+
}
223+
"i" => {
224+
// c_int
225+
const SIZE: usize = mem::size_of::<c_int>();
226+
let val = if bytes.len() >= SIZE {
227+
c_int::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
228+
} else {
229+
0
230+
};
231+
Ok(vm.ctx.new_int(val).into())
232+
}
233+
"I" => {
234+
// c_uint
235+
const SIZE: usize = mem::size_of::<c_uint>();
236+
let val = if bytes.len() >= SIZE {
237+
c_uint::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
238+
} else {
239+
0
240+
};
241+
Ok(vm.ctx.new_int(val).into())
242+
}
243+
"l" => {
244+
// c_long
245+
const SIZE: usize = mem::size_of::<c_long>();
246+
let val = if bytes.len() >= SIZE {
247+
c_long::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
248+
} else {
249+
0
250+
};
251+
Ok(vm.ctx.new_int(val).into())
252+
}
253+
"L" => {
254+
// c_ulong
255+
const SIZE: usize = mem::size_of::<c_ulong>();
256+
let val = if bytes.len() >= SIZE {
257+
c_ulong::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
258+
} else {
259+
0
260+
};
261+
Ok(vm.ctx.new_int(val).into())
262+
}
263+
"q" => {
264+
// c_longlong
265+
const SIZE: usize = mem::size_of::<c_longlong>();
266+
let val = if bytes.len() >= SIZE {
267+
c_longlong::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
268+
} else {
269+
0
270+
};
271+
Ok(vm.ctx.new_int(val).into())
272+
}
273+
"Q" => {
274+
// c_ulonglong
275+
const SIZE: usize = mem::size_of::<c_ulonglong>();
276+
let val = if bytes.len() >= SIZE {
277+
c_ulonglong::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
278+
} else {
279+
0
280+
};
281+
Ok(vm.ctx.new_int(val).into())
282+
}
283+
"f" => {
284+
// c_float
285+
const SIZE: usize = mem::size_of::<c_float>();
286+
let val = if bytes.len() >= SIZE {
287+
c_float::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
288+
} else {
289+
0.0
290+
};
291+
Ok(vm.ctx.new_float(val as f64).into())
292+
}
293+
"d" | "g" => {
294+
// c_double
295+
const SIZE: usize = mem::size_of::<c_double>();
296+
let val = if bytes.len() >= SIZE {
297+
c_double::from_ne_bytes(bytes[..SIZE].try_into().expect("size checked"))
298+
} else {
299+
0.0
300+
};
301+
Ok(vm.ctx.new_float(val).into())
302+
}
303+
"?" => {
304+
// c_bool
305+
let val = !bytes.is_empty() && bytes[0] != 0;
306+
Ok(vm.ctx.new_bool(val).into())
307+
}
308+
"P" | "z" | "Z" => {
309+
// Pointer types - return as integer address
310+
let val = if bytes.len() >= mem::size_of::<libc::uintptr_t>() {
311+
const UINTPTR_LEN: usize = mem::size_of::<libc::uintptr_t>();
312+
let mut arr = [0u8; UINTPTR_LEN];
313+
arr[..bytes.len().min(UINTPTR_LEN)]
314+
.copy_from_slice(&bytes[..bytes.len().min(UINTPTR_LEN)]);
315+
usize::from_ne_bytes(arr)
316+
} else {
317+
0
318+
};
319+
Ok(vm.ctx.new_int(val).into())
320+
}
321+
"u" => {
322+
// c_wchar - wide character
323+
let val = if bytes.len() >= mem::size_of::<WideChar>() {
324+
let wc = if mem::size_of::<WideChar>() == 2 {
325+
u16::from_ne_bytes([bytes[0], bytes[1]]) as u32
223326
} else {
224-
0
225-
};
226-
Ok(vm.ctx.new_int(val).into())
227-
}
228-
"I" | "L" => {
229-
// c_uint, c_ulong (assuming 32-bit)
230-
let val = if bytes.len() >= 4 {
231327
u32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
232-
} else {
233-
0
234328
};
235-
Ok(vm.ctx.new_int(val).into())
236-
}
237-
"q" => {
238-
// c_longlong
239-
let val = if bytes.len() >= 8 {
240-
i64::from_ne_bytes([
241-
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5],
242-
bytes[6], bytes[7],
243-
])
244-
} else {
245-
0
246-
};
247-
Ok(vm.ctx.new_int(val).into())
248-
}
249-
"Q" => {
250-
// c_ulonglong
251-
let val = if bytes.len() >= 8 {
252-
u64::from_ne_bytes([
253-
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5],
254-
bytes[6], bytes[7],
255-
])
256-
} else {
257-
0
258-
};
259-
Ok(vm.ctx.new_int(val).into())
260-
}
261-
"f" => {
262-
// c_float
263-
let val = if bytes.len() >= 4 {
264-
f32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
265-
} else {
266-
0.0
267-
};
268-
Ok(vm.ctx.new_float(val as f64).into())
269-
}
270-
"d" | "g" => {
271-
// c_double
272-
let val = if bytes.len() >= 8 {
273-
f64::from_ne_bytes([
274-
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5],
275-
bytes[6], bytes[7],
276-
])
277-
} else {
278-
0.0
279-
};
280-
Ok(vm.ctx.new_float(val).into())
281-
}
282-
"?" => {
283-
// c_bool
284-
let val = !bytes.is_empty() && bytes[0] != 0;
285-
Ok(vm.ctx.new_bool(val).into())
286-
}
287-
"P" | "z" | "Z" => {
288-
// Pointer types - return as integer address
289-
let val = if bytes.len() >= mem::size_of::<usize>() {
290-
let mut arr = [0u8; 8];
291-
arr[..bytes.len().min(8)].copy_from_slice(&bytes[..bytes.len().min(8)]);
292-
usize::from_ne_bytes(arr)
293-
} else {
294-
0
295-
};
296-
Ok(vm.ctx.new_int(val).into())
297-
}
298-
"u" => {
299-
// c_wchar - wide character
300-
let val = if bytes.len() >= mem::size_of::<WideChar>() {
301-
let wc = if mem::size_of::<WideChar>() == 2 {
302-
u16::from_ne_bytes([bytes[0], bytes[1]]) as u32
303-
} else {
304-
u32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
305-
};
306-
char::from_u32(wc).unwrap_or('\0')
307-
} else {
308-
'\0'
309-
};
310-
Ok(vm.ctx.new_str(val.to_string()).into())
311-
}
312-
_ => Ok(vm.ctx.none()),
313-
};
314-
}
329+
char::from_u32(wc).unwrap_or('\0')
330+
} else {
331+
'\0'
332+
};
333+
Ok(vm.ctx.new_str(val.to_string()).into())
334+
}
335+
_ => Ok(vm.ctx.none()),
336+
};
315337
}
316338
// Default: return bytes as-is
317339
Ok(vm.ctx.new_bytes(bytes.to_vec()).into())
@@ -339,9 +361,11 @@ pub(crate) mod _ctypes {
339361
} else if !SIMPLE_TYPE_CHARS.contains(tp_str.as_str()) {
340362
Err(vm.new_attribute_error(format!("class must define a '_type_' attribute which must be\n a single character string containing one of {SIMPLE_TYPE_CHARS}, currently it is {tp_str}.")))
341363
} else {
364+
let size = get_size(&tp_str);
342365
Ok(PyCSimple {
343366
_type_: tp_str,
344367
value: AtomicCell::new(vm.ctx.none()),
368+
cdata: rustpython_common::lock::PyRwLock::new(CDataObject::new(size)),
345369
})
346370
}
347371
} else {

0 commit comments

Comments
 (0)