#![feature(doc_cfg, external_doc)]
#![doc(include = "../README.md")]
#![cfg_attr(feature = "extension-module", crate_type = "cdylib")]
extern crate pyo3;
#[cfg(feature = "extension-module")]
extern crate pyo3_built;
#[cfg(feature = "extension-module")]
mod built;
use std::ops::Deref;
use pyo3::class::basic::CompareOp;
use pyo3::class::PyGCProtocol;
use pyo3::class::PyIterProtocol;
use pyo3::class::PyNumberProtocol;
use pyo3::class::PyObjectProtocol;
use pyo3::class::PySequenceProtocol;
use pyo3::exceptions::KeyError;
use pyo3::gc::PyTraverseError;
use pyo3::gc::PyVisit;
use pyo3::prelude::*;
use pyo3::types::PyAny;
use pyo3::types::PyDict;
use pyo3::types::PyFrozenSet;
use pyo3::types::PyIterator;
use pyo3::types::PySet;
use pyo3::types::PyTuple;
use pyo3::AsPyPointer;
use pyo3::PyNativeType;
macro_rules! common_impl {
($cls:ty) => {
impl $cls {
pub fn new() -> Self {
Self::default()
}
pub fn from_set(obj: PyObject) -> Self {
Self { inner: Some(obj) }
}
pub fn try_from_any(py: Python, any: &PyAny) -> PyResult<Self> {
Self::try_from_obj(py, any.to_object(py))
}
pub fn try_from_obj(py: Python, obj: PyObject) -> PyResult<Self> {
if let Ok(s) = obj.cast_as::<PySet>(py) {
if s.is_empty() {
Ok(Self::new())
} else {
s.to_object(py).call_method0(py, "copy").map(Self::from_set)
}
} else if let Ok(d) = obj.cast_as::<PyDict>(py) {
if d.is_empty() {
Ok(Self::new())
} else {
unsafe {
let set = pyo3::ffi::PySet_New(d.as_ptr());
Ok(Self::from_set(PyObject::from_owned_ptr(py, set)))
}
}
} else {
let iterator = PyIterator::from_object(py, &obj)?;
Self::try_from_iterator(py, iterator)
}
}
pub fn try_from_iterator(py: Python, it: PyIterator) -> PyResult<Self> {
let items: PyResult<Vec<&PyAny>> = it.collect();
let res = items?;
if res.is_empty() {
Ok(Self::default())
} else {
let set = PySet::new(py, res.as_slice())?;
Ok(Self::from_set(set.to_object(py)))
}
}
}
impl FromPy<PySet> for $cls {
fn from_py(set: PySet, py: Python) -> Self {
Self::from_set(set.to_object(py))
}
}
impl Drop for $cls {
fn drop(&mut self) {
if let Some(set) = self.inner.take() {
Python::acquire_gil().python().release(set)
}
}
}
#[pymethods]
impl $cls {
#[new]
fn __new__(obj: &PyRawObject, iterable: Option<&PyAny>) -> PyResult<()> {
let mut val = Self::new();
val.__init__(iterable)?;
obj.init(val);
Ok(())
}
fn __init__(&mut self, iterable: Option<&PyAny>) -> PyResult<()> {
if let Some(it) = iterable {
let gil = Python::acquire_gil();
let py = gil.python();
*self = Self::try_from_any(py, it)?;
} else {
self.inner = None;
}
Ok(())
}
fn __reduce__(&self) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
let ty = <$cls as pyo3::type_object::PyTypeObject>::type_object();
match self.inner {
None => Ok((ty, PyTuple::empty(py)).to_object(py)),
Some(ref set) => Ok((ty, (set.call_method0(py, "copy")?,)).to_object(py)),
}
}
fn add(&mut self, item: &PyAny) -> PyResult<()> {
let gil = Python::acquire_gil();
let py = gil.python();
if self.inner.is_none() {
self.inner = Some(PySet::empty(py)?.to_object(py));
}
self.inner.as_ref().unwrap().cast_as::<PySet>(py)?.add(item)
}
fn clear(&mut self) -> PyResult<()> {
if let Some(inner) = self.inner.take() {
Python::acquire_gil().python().release(inner)
};
Ok(())
}
fn copy(&self) -> PyResult<Self> {
match self.inner {
None => Ok(Self::new()),
Some(ref inner) => {
let gil = Python::acquire_gil();
let py = gil.python();
inner.call_method0(py, "copy").map(Self::from_set)
}
}
}
#[args(others = "*")]
fn difference(&self, others: &PyTuple) -> PyResult<Self> {
let py = others.py();
if let Some(ref obj) = self.inner {
obj.call_method1(py, "difference", others)
.and_then(|s| Self::try_from_obj(py, s))
} else {
for arg in others.iter() {
if let Err(e) = PyIterator::from_object(py, arg) {
return Err(e.into());
}
}
Ok(Self::new())
}
}
#[args(others = "*")]
fn difference_update(&mut self, others: &PyTuple) -> PyResult<()> {
let gil = Python::acquire_gil();
let py = gil.python();
if let Some(ref inner) = self.inner {
inner.call_method1(py, "difference_update", others)?;
if inner.cast_as::<PySet>(py).unwrap().is_empty() {
self.inner = None;
}
} else {
for arg in others.iter() {
if let Err(e) = PyIterator::from_object(py, arg) {
return Err(e.into());
}
}
}
Ok(())
}
fn discard(&mut self, elem: &PyAny) -> PyResult<()> {
if let Some(ref set) = self.inner {
let gil = Python::acquire_gil();
set.call_method1(gil.python(), "discard", (elem,))?;
}
Ok(())
}
#[args(others = "*")]
fn intersection(&self, others: &PyTuple) -> PyResult<Self> {
let py = others.py();
if let Some(ref obj) = self.inner {
obj.call_method1(py, "intersection", others)
.and_then(|s| Self::try_from_obj(py, s))
} else {
Ok(Self::new())
}
}
#[args(others = "*")]
fn intersection_update(&mut self, others: &PyTuple) -> PyResult<()> {
if let Some(ref inner) = self.inner {
let gil = Python::acquire_gil();
let py = gil.python();
inner.call_method1(py, "intersection_update", others)?;
if inner.cast_as::<PySet>(py).unwrap().is_empty() {
self.inner = None;
}
}
Ok(())
}
fn isdisjoint(&self, other: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
match self.inner {
None => Ok(true.to_object(py)),
Some(ref obj) => obj.call_method1(py, "isdisjoint", (other,)),
}
}
fn issubset(&self, other: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
match self.inner {
None => Ok(true.to_object(py)),
Some(ref obj) => obj.call_method1(py, "issubset", (other,)),
}
}
fn issuperset(&self, other: &PyAny) -> PyResult<bool> {
let gil = Python::acquire_gil();
let py = gil.python();
let mut it = PyIterator::from_object(py, other)?;
match self.inner {
None => Ok(it.next().is_none()),
Some(ref inner) => {
let set = inner.cast_as::<PySet>(py).unwrap();
for item in it {
if !set.contains(item?)? {
return Ok(false);
}
}
Ok(true)
}
}
}
fn pop(&mut self) -> PyResult<PyObject> {
if let Some(ref inner) = self.inner {
let gil = Python::acquire_gil();
let set = inner.cast_as::<PySet>(gil.python())?;
if let Some(item) = set.pop() {
if set.len() == 0 {
self.inner = None;
}
return Ok(item);
}
}
KeyError::into("pop from an empty set")
}
fn remove(&mut self, item: &PyAny) -> PyResult<()> {
let gil = Python::acquire_gil();
let py = gil.python();
if let Some(ref inner) = self.inner {
if let Ok(ref other) = item.cast_as::<$cls>() {
if let Some(ref obj) = other.inner {
inner.call_method1(py, "remove", (obj.clone_ref(py),))?
} else {
inner.call_method1(py, "remove", (PyFrozenSet::empty(py)?,))?
};
} else {
inner.call_method1(py, "remove", (item,))?;
}
if inner.cast_as::<PySet>(py).unwrap().is_empty() {
self.inner = None
}
Ok(())
} else {
KeyError::into(item.to_object(py))
}
}
fn symmetric_difference(&self, other: &PyAny) -> PyResult<Self> {
let gil = Python::acquire_gil();
let py = gil.python();
if let Some(ref obj) = self.inner {
obj.call_method1(py, "symmetric_difference", (other,))
.and_then(|obj| Self::try_from_obj(py, obj))
} else {
Self::try_from_any(py, other)
}
}
#[args(others = "*")]
fn symmetric_difference_update(&self, others: &PyTuple) -> PyResult<()> {
if let Some(ref inner) = self.inner {
let gil = Python::acquire_gil();
let py = gil.python();
inner.call_method1(py, "symmetric_difference_update", others)?;
}
Ok(())
}
#[args(others = "*")]
fn union(&self, others: &PyTuple) -> PyResult<Self> {
let py = others.py();
if let Some(ref inner) = self.inner {
inner.call_method1(py, "union", others).map(Self::from_set)
} else {
let s = PySet::empty(py)?.to_object(py);
s.call_method1(py, "union", others)
.and_then(|set| Self::try_from_obj(py, set))
}
}
#[args(others = "*")]
fn update(&mut self, others: &PyTuple) -> PyResult<()> {
for other in others.iter() {
match self.inner {
None => {
*self = Self::try_from_any(others.py(), other)?;
}
Some(ref inner) => {
inner.call_method1(others.py(), "update", (other,))?;
}
}
}
Ok(())
}
}
#[pyproto]
impl PyIterProtocol for $cls {
fn __iter__(slf: PyRefMut<'p, Self>) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
match slf.deref().inner {
None => PyTuple::empty(py)
.to_object(py)
.call_method0(py, "__iter__"),
Some(ref inner) => inner.call_method0(py, "__iter__"),
}
}
}
#[pyproto]
impl PyNumberProtocol for $cls {
fn __and__(lhs: &Self, rhs: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
if rhs.cast_as::<PySet>().is_err()
&& rhs.cast_as::<PyFrozenSet>().is_err()
&& rhs.cast_as::<$cls>().is_err()
{
return Ok(py.NotImplemented());
}
let args: Py<PyTuple> = (rhs,).into_py(gil.python());
lhs.intersection(&args.as_ref(py))
.and_then(|s| Py::new(py, s))
.map(PyObject::from)
}
fn __sub__(lhs: &Self, rhs: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
if rhs.cast_as::<PySet>().is_err()
&& rhs.cast_as::<PyFrozenSet>().is_err()
&& rhs.cast_as::<$cls>().is_err()
{
return Ok(py.NotImplemented());
}
let args: Py<PyTuple> = (rhs,).into_py(gil.python());
lhs.difference(&args.as_ref(py))
.and_then(|s| Py::new(py, s))
.map(PyObject::from)
}
fn __or__(lhs: &Self, rhs: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
if rhs.cast_as::<PySet>().is_err()
&& rhs.cast_as::<PyFrozenSet>().is_err()
&& rhs.cast_as::<$cls>().is_err()
{
return Ok(py.NotImplemented());
}
let args: Py<PyTuple> = (rhs,).into_py(gil.python());
lhs.union(&args.as_ref(py))
.and_then(|s| Py::new(py, s))
.map(PyObject::from)
}
fn __xor__(lhs: &Self, rhs: &PyAny) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
if rhs.cast_as::<PySet>().is_err()
&& rhs.cast_as::<PyFrozenSet>().is_err()
&& rhs.cast_as::<$cls>().is_err()
{
return Ok(py.NotImplemented());
}
lhs.symmetric_difference(rhs)
.and_then(|s| Py::new(py, s))
.map(PyObject::from)
}
}
#[pyproto]
impl PyObjectProtocol for $cls {
fn __repr__(&self) -> PyResult<PyObject> {
let gil = Python::acquire_gil();
let py = gil.python();
match self.inner {
None => {
Ok("set()".to_object(py))
}
Some(ref inner) => {
inner.call_method0(py, "__repr__")
}
}
}
fn __bool__(&self) -> PyResult<bool> {
Ok(self.inner.is_some())
}
fn __richcmp__(&self, obj: &PyAny, op: CompareOp) -> PyResult<PyObject> {
use self::CompareOp::*;
let gil = Python::acquire_gil();
let py = gil.python();
if let Ok(other) = obj.cast_as::<Self>() {
match (&self.inner, &other.inner) {
(None, None) => match op {
Eq | Le | Ge => Ok(true.to_object(py)),
Ne | Lt | Gt => Ok(false.to_object(py)),
},
(None, Some(_)) => match op {
Eq | Gt | Ge => Ok(false.to_object(py)),
Ne | Lt | Le => Ok(true.to_object(py)),
},
(Some(_), None) => match op {
Eq | Lt | Le => Ok(false.to_object(py)),
Ne | Gt | Ge => Ok(true.to_object(py)),
},
(Some(l), Some(r)) => match op {
Eq => l.call_method1(py, "__eq__", (r,)),
Ne => l.call_method1(py, "__ne__", (r,)),
Lt => l.call_method1(py, "__lt__", (r,)),
Le => l.call_method1(py, "__le__", (r,)),
Gt => l.call_method1(py, "__gt__", (r,)),
Ge => l.call_method1(py, "__ge__", (r,)),
},
}
} else if let Ok(other) = obj.cast_as::<PySet>() {
match (&self.inner, other.is_empty()) {
(None, true) => match op {
Eq | Le | Ge => Ok(true.to_object(py)),
Ne | Lt | Gt => Ok(false.to_object(py)),
},
(None, false) => match op {
Eq | Gt | Ge => Ok(false.to_object(py)),
Ne | Lt | Le => Ok(true.to_object(py)),
},
(Some(l), _) => match op {
Eq => l.call_method1(py, "__eq__", (obj,)),
Ne => l.call_method1(py, "__ne__", (obj,)),
Lt => l.call_method1(py, "__lt__", (obj,)),
Le => l.call_method1(py, "__le__", (obj,)),
Gt => l.call_method1(py, "__gt__", (obj,)),
Ge => l.call_method1(py, "__ge__", (obj,)),
},
}
} else if let Ok(other) = obj.cast_as::<PyFrozenSet>() {
match (&self.inner, other.is_empty()) {
(None, true) => match op {
Eq | Le | Ge => Ok(true.to_object(py)),
Ne | Lt | Gt => Ok(false.to_object(py)),
},
(None, false) => match op {
Eq | Gt | Ge => Ok(false.to_object(py)),
Ne | Lt | Le => Ok(true.to_object(py)),
},
(Some(l), _) => match op {
Eq => l.call_method1(py, "__eq__", (obj,)),
Ne => l.call_method1(py, "__ne__", (obj,)),
Lt => l.call_method1(py, "__lt__", (obj,)),
Le => l.call_method1(py, "__le__", (obj,)),
Gt => l.call_method1(py, "__gt__", (obj,)),
Ge => l.call_method1(py, "__ge__", (obj,)),
},
}
} else {
match op {
CompareOp::Eq => Ok(false.to_object(py)),
CompareOp::Ne => Ok(true.to_object(py)),
_ => Ok(py.NotImplemented()),
}
}
}
}
#[pyproto]
impl PySequenceProtocol for $cls {
fn __len__(&self) -> PyResult<usize> {
match self.inner {
None => Ok(0usize),
Some(ref inner) => {
let gil = Python::acquire_gil();
Ok(inner.cast_as::<PySet>(gil.python())?.len())
}
}
}
fn __contains__(&self, item: &PyAny) -> PyResult<bool> {
let gil = Python::acquire_gil();
let py = gil.python();
if let Some(ref obj) = self.inner {
if let Ok(ref other) = item.cast_as::<$cls>() {
let set = obj.cast_as::<PySet>(py).unwrap();
match other.inner {
Some(ref obj) => set.contains(obj),
None => set.contains(PyFrozenSet::empty(py)?),
}
} else {
obj.call_method1(py, "__contains__", (item,))
.and_then(|val| val.extract(py))
}
} else {
Ok(false)
}
}
}
};
}
#[pyclass(gc, module = "nanoset")]
#[derive(Debug, Default)]
pub struct NanoSet {
inner: Option<PyObject>,
}
common_impl!(NanoSet);
#[pyproto]
impl PyGCProtocol for NanoSet {
fn __traverse__(&'p self, visit: PyVisit) -> Result<(), PyTraverseError> {
match self.inner {
Some(ref obj) => visit.call(obj),
None => Ok(()),
}
}
fn __clear__(&'p mut self) {
if let Some(obj) = self.inner.take() {
let gil = Python::acquire_gil();
gil.python().release(obj)
}
}
}
#[pyclass(module = "nanoset")]
#[derive(Debug, Default)]
pub struct PicoSet {
inner: Option<PyObject>,
}
common_impl!(PicoSet);
#[cfg(feature = "extension-module")]
#[pymodule]
fn nanoset(py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<NanoSet>()?;
m.add_class::<PicoSet>()?;
m.add("__build__", pyo3_built::pyo3_built!(py, built))?;
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
m.add("__author__", env!("CARGO_PKG_AUTHORS").replace(':', "\n"))?;
let cabc = py.import("collections.abc")?;
let set = cabc.get("Set")?.to_object(py);
set.call_method1(
py,
"register",
(<NanoSet as pyo3::type_object::PyTypeObject>::type_object(),),
)?;
set.call_method1(
py,
"register",
(<PicoSet as pyo3::type_object::PyTypeObject>::type_object(),),
)?;
Ok(())
}