Skip to content

Commit a136f90

Browse files
Add basic dict function to c-api (RustPython#7929)
* Add basic dict function to c-api * Fix iter * Do not use mapping protocol
1 parent 948924a commit a136f90

3 files changed

Lines changed: 146 additions & 1 deletion

File tree

crates/capi/src/dictobject.rs

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
use crate::PyObject;
2+
use crate::object::define_py_check;
3+
use crate::pystate::with_vm;
4+
use core::ffi::c_int;
5+
use core::ptr::NonNull;
6+
use rustpython_vm::AsObject;
7+
use rustpython_vm::builtins::PyDict;
8+
9+
define_py_check!(fn PyDict_Check, types.dict_type);
10+
define_py_check!(exact fn PyDict_CheckExact, types.dict_type);
11+
define_py_check!(fn PyDictKeys_Check, types.dict_keys_type);
12+
define_py_check!(fn PyDictValues_Check, types.dict_values_type);
13+
define_py_check!(fn PyDictItems_Check, types.dict_items_type);
14+
15+
#[unsafe(no_mangle)]
16+
pub extern "C" fn PyDict_New() -> *mut PyObject {
17+
with_vm(|vm| vm.ctx.new_dict())
18+
}
19+
20+
#[unsafe(no_mangle)]
21+
pub unsafe extern "C" fn PyDict_SetItem(
22+
dict: *mut PyObject,
23+
key: *mut PyObject,
24+
val: *mut PyObject,
25+
) -> c_int {
26+
with_vm(|vm| {
27+
let dict = unsafe { &*dict }.try_downcast_ref::<PyDict>(vm)?;
28+
let key = unsafe { &*key };
29+
let value = unsafe { &*val }.to_owned();
30+
dict.inner_setitem(key, value, vm)
31+
})
32+
}
33+
34+
#[unsafe(no_mangle)]
35+
pub unsafe extern "C" fn PyDict_GetItemRef(
36+
dict: *mut PyObject,
37+
key: *mut PyObject,
38+
result: *mut *mut PyObject,
39+
) -> c_int {
40+
with_vm(|vm| {
41+
unsafe { *result = core::ptr::null_mut() };
42+
let dict = unsafe { &*dict }.try_downcast_ref::<PyDict>(vm)?;
43+
let key = unsafe { &*key };
44+
45+
if let Some(value) = dict.inner_getitem_opt(key, vm)? {
46+
unsafe {
47+
*result = value.into_raw().as_ptr();
48+
}
49+
Ok(true)
50+
} else {
51+
unsafe {
52+
*result = core::ptr::null_mut();
53+
}
54+
Ok(false)
55+
}
56+
})
57+
}
58+
59+
#[unsafe(no_mangle)]
60+
pub unsafe extern "C" fn PyDict_Size(dict: *mut PyObject) -> isize {
61+
with_vm(|vm| {
62+
let dict = unsafe { &*dict }.try_downcast_ref::<PyDict>(vm)?;
63+
Ok(dict.__len__())
64+
})
65+
}
66+
67+
#[unsafe(no_mangle)]
68+
pub unsafe extern "C" fn PyDict_Next(
69+
dict: *mut PyObject,
70+
pos: *mut isize,
71+
key: *mut *mut PyObject,
72+
value: *mut *mut PyObject,
73+
) -> c_int {
74+
with_vm(|vm| {
75+
let dict = unsafe { &*dict }.try_downcast_ref::<PyDict>(vm)?;
76+
let index = unsafe { *pos } as usize;
77+
78+
if let Some((next_pos, k, v)) = dict.next_entry(index) {
79+
unsafe {
80+
*pos = next_pos as isize;
81+
if let Some(key) = NonNull::new(key) {
82+
key.write(k.as_object().as_raw().cast_mut());
83+
}
84+
if let Some(value) = NonNull::new(value) {
85+
value.write(v.as_object().as_raw().cast_mut());
86+
}
87+
}
88+
Ok(true)
89+
} else {
90+
Ok(false)
91+
}
92+
})
93+
}
94+
95+
#[cfg(false)]
96+
mod tests {
97+
use pyo3::prelude::*;
98+
use pyo3::types::{IntoPyDict, PyDict, PyInt};
99+
100+
#[test]
101+
fn test_create_empty_dict() {
102+
Python::attach(|py| {
103+
let dict = PyDict::new(py);
104+
assert!(dict.is_instance_of::<PyDict>());
105+
})
106+
}
107+
108+
#[test]
109+
fn test_create_dict_with_items() {
110+
Python::attach(|py| {
111+
let dict = [(1, 2), (3, 4)].into_py_dict(py)?;
112+
let value = dict.get_item(1)?.unwrap().cast_into::<PyInt>()?;
113+
assert_eq!(value, 2);
114+
assert_eq!(dict.len(), 2);
115+
116+
Ok::<_, PyErr>(())
117+
})
118+
.unwrap()
119+
}
120+
121+
#[test]
122+
fn test_dict_iter() {
123+
Python::attach(|py| {
124+
let dict = [(1, 2), (3, 4)].into_py_dict(py).unwrap();
125+
let values = dict
126+
.into_iter()
127+
.flat_map(|(k, v)| [k.extract().unwrap(), v.extract().unwrap()])
128+
.collect::<Vec<u32>>();
129+
assert_eq!(values, vec![1, 2, 3, 4]);
130+
})
131+
}
132+
}

crates/capi/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ extern crate alloc;
1111
pub mod abstract_;
1212
pub mod bytesobject;
1313
pub mod ceval;
14+
pub mod dictobject;
1415
pub mod import;
1516
pub mod longobject;
1617
pub mod object;

crates/vm/src/builtins/dict.rs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ impl PyDict {
205205

206206
/// Set item variant which can be called with multiple
207207
/// key types, such as str to name a notable one.
208-
pub(crate) fn inner_setitem<K: DictKey + ?Sized>(
208+
pub fn inner_setitem<K: DictKey + ?Sized>(
209209
&self,
210210
key: &K,
211211
value: PyObjectRef,
@@ -248,6 +248,18 @@ impl PyDict {
248248
pub fn size(&self) -> dict_inner::DictSize {
249249
self.entries.size()
250250
}
251+
252+
pub fn next_entry(&self, position: usize) -> Option<(usize, PyObjectRef, PyObjectRef)> {
253+
self.entries.next_entry(position)
254+
}
255+
256+
pub fn inner_getitem_opt<K: DictKey + ?Sized>(
257+
&self,
258+
key: &K,
259+
vm: &VirtualMachine,
260+
) -> PyResult<Option<PyObjectRef>> {
261+
self.entries.get(vm, key)
262+
}
251263
}
252264

253265
// Python dict methods:

0 commit comments

Comments
 (0)