Skip to content

Commit 885cf5c

Browse files
bschoenmaeckersbschoenmaeckers
authored
Add set object functions to c-api (RustPython#8002)
1 parent 3e1c3bc commit 885cf5c

5 files changed

Lines changed: 170 additions & 5 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/capi/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ crate-type = ["cdylib", "rlib"]
1313

1414
[dependencies]
1515
bitflags = { workspace = true }
16+
itertools = { workspace = true }
1617
num-complex = { workspace = true }
1718
rustpython-vm = { workspace = true, features = ["threading", "compiler"] }
1819
rustpython-stdlib = {workspace = true, features = ["threading"] }

‎crates/capi/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ pub mod pyerrors;
2525
pub mod pylifecycle;
2626
pub mod pystate;
2727
pub mod refcount;
28+
pub mod setobject;
2829
pub mod traceback;
2930
pub mod tupleobject;
3031
pub mod unicodeobject;

‎crates/capi/src/setobject.rs‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
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 itertools::process_results;
6+
use rustpython_vm::AsObject;
7+
use rustpython_vm::PyPayload;
8+
use rustpython_vm::TryFromObject;
9+
use rustpython_vm::builtins::{PyFrozenSet, PySet};
10+
use rustpython_vm::function::ArgIterable;
11+
12+
define_py_check!(fn PySet_Check, types.set_type);
13+
define_py_check!(fn PyFrozenSet_Check, types.frozenset_type);
14+
15+
#[unsafe(no_mangle)]
16+
pub unsafe extern "C" fn PySet_New(iterable: *mut PyObject) -> *mut PyObject {
17+
with_vm(|vm| {
18+
if iterable.is_null() {
19+
return Ok(PySet::default().into_ref(&vm.ctx));
20+
}
21+
22+
let iterable = ArgIterable::try_from_object(vm, unsafe { &*iterable }.to_owned())?;
23+
let set = PySet::default().into_ref(&vm.ctx);
24+
for item in iterable.iter(vm)? {
25+
set.add(item?, vm)?;
26+
}
27+
Ok(set)
28+
})
29+
}
30+
31+
#[unsafe(no_mangle)]
32+
pub unsafe extern "C" fn PyFrozenSet_New(iterable: *mut PyObject) -> *mut PyObject {
33+
with_vm(|vm| {
34+
if iterable.is_null() {
35+
return Ok(vm.ctx.empty_frozenset.to_owned());
36+
}
37+
38+
let iterable = ArgIterable::try_from_object(vm, unsafe { &*iterable }.to_owned())?;
39+
let set = process_results(iterable.iter(vm)?, |it| PyFrozenSet::from_iter(vm, it))??;
40+
Ok(set.into_ref(&vm.ctx))
41+
})
42+
}
43+
44+
#[unsafe(no_mangle)]
45+
pub unsafe extern "C" fn PySet_Add(set: *mut PyObject, key: *mut PyObject) -> c_int {
46+
with_vm(|vm| {
47+
let set = unsafe { &*set }.try_downcast_ref::<PySet>(vm)?;
48+
let key = unsafe { &*key }.to_owned();
49+
set.add(key, vm)
50+
})
51+
}
52+
53+
#[unsafe(no_mangle)]
54+
pub unsafe extern "C" fn PySet_Clear(set: *mut PyObject) -> c_int {
55+
with_vm(|vm| {
56+
let set = unsafe { &*set }.try_downcast_ref::<PySet>(vm)?;
57+
set.clear();
58+
Ok(())
59+
})
60+
}
61+
62+
#[unsafe(no_mangle)]
63+
pub unsafe extern "C" fn PySet_Contains(anyset: *mut PyObject, key: *mut PyObject) -> c_int {
64+
with_vm(|vm| {
65+
let anyset = unsafe { &*anyset };
66+
let key = unsafe { &*key };
67+
68+
if let Some(set) = anyset.downcast_ref::<PySet>() {
69+
set.__contains__(key, vm)
70+
} else if let Some(frozenset) = anyset.downcast_ref::<PyFrozenSet>() {
71+
frozenset.__contains__(key, vm)
72+
} else {
73+
Err(vm.new_type_error(format!(
74+
"expected set or frozenset, got '{}'",
75+
anyset.class().name()
76+
)))
77+
}
78+
})
79+
}
80+
81+
#[unsafe(no_mangle)]
82+
pub unsafe extern "C" fn PySet_Discard(set: *mut PyObject, key: *mut PyObject) -> c_int {
83+
with_vm(|vm| {
84+
let set = unsafe { &*set }.try_downcast_ref::<PySet>(vm)?;
85+
let key = unsafe { &*key };
86+
let had_item = set.__contains__(key, vm)?;
87+
if had_item {
88+
set.discard(key.to_owned(), vm)?;
89+
}
90+
Ok(had_item)
91+
})
92+
}
93+
94+
#[unsafe(no_mangle)]
95+
pub unsafe extern "C" fn PySet_Pop(set: *mut PyObject) -> *mut PyObject {
96+
with_vm(|vm| {
97+
let set = unsafe { &*set }.try_downcast_ref::<PySet>(vm)?;
98+
set.pop(vm)
99+
})
100+
}
101+
102+
#[unsafe(no_mangle)]
103+
pub unsafe extern "C" fn PySet_Size(anyset: *mut PyObject) -> isize {
104+
with_vm(|vm| {
105+
let anyset = unsafe { &*anyset };
106+
if let Some(set) = anyset.downcast_ref::<PySet>() {
107+
set.as_object().length(vm)
108+
} else if let Some(frozenset) = anyset.downcast_ref::<PyFrozenSet>() {
109+
frozenset.as_object().length(vm)
110+
} else {
111+
Err(vm.new_type_error(format!(
112+
"expected set or frozenset, got '{}'",
113+
anyset.class().name()
114+
)))
115+
}
116+
})
117+
}
118+
119+
#[cfg(false)]
120+
mod tests {
121+
use pyo3::prelude::*;
122+
use pyo3::types::{PyFrozenSet, PyInt, PySet};
123+
124+
#[test]
125+
fn new_and_size() {
126+
Python::attach(|py| {
127+
let set = PySet::empty(py).unwrap();
128+
assert!(set.is_instance_of::<PySet>());
129+
assert_eq!(set.len(), 0);
130+
131+
let frozen = PyFrozenSet::empty(py).unwrap();
132+
assert!(frozen.is_instance_of::<PyFrozenSet>());
133+
assert_eq!(frozen.len(), 0);
134+
})
135+
}
136+
137+
#[test]
138+
fn add_contains_discard() {
139+
Python::attach(|py| {
140+
let set = PySet::empty(py).unwrap();
141+
let item = PyInt::new(py, 42);
142+
143+
set.add(&item).unwrap();
144+
assert!(set.contains(&item).unwrap());
145+
set.discard(&item).unwrap();
146+
assert!(!set.contains(&item).unwrap());
147+
})
148+
}
149+
150+
#[test]
151+
fn pop_reduces_size() {
152+
Python::attach(|py| {
153+
let set = PySet::empty(py).unwrap();
154+
set.add(7).unwrap();
155+
assert_eq!(set.len(), 1);
156+
157+
let popped = set.pop().unwrap();
158+
assert_eq!(popped.extract::<i32>().unwrap(), 7);
159+
assert_eq!(set.len(), 0);
160+
})
161+
}
162+
}

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -537,7 +537,7 @@ impl PySet {
537537
self.inner.len()
538538
}
539539

540-
fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
540+
pub fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
541541
self.inner.contains(needle, vm)
542542
}
543543

@@ -679,17 +679,17 @@ impl PySet {
679679
}
680680

681681
#[pymethod]
682-
fn discard(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
682+
pub fn discard(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
683683
self.inner.discard(&item, vm).map(|_| ())
684684
}
685685

686686
#[pymethod]
687-
fn clear(&self) {
687+
pub fn clear(&self) {
688688
self.inner.clear()
689689
}
690690

691691
#[pymethod]
692-
fn pop(&self, vm: &VirtualMachine) -> PyResult {
692+
pub fn pop(&self, vm: &VirtualMachine) -> PyResult {
693693
self.inner.pop(vm)
694694
}
695695

@@ -995,7 +995,7 @@ impl PyFrozenSet {
995995
self.inner.len()
996996
}
997997

998-
fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
998+
pub fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
999999
self.inner.contains(needle, vm)
10001000
}
10011001

0 commit comments

Comments
 (0)