pyo3/types/
mapping.rs

1use crate::conversion::IntoPyObject;
2use crate::err::PyResult;
3use crate::ffi_ptr_ext::FfiPtrExt;
4use crate::instance::Bound;
5use crate::py_result_ext::PyResultExt;
6use crate::sync::PyOnceLock;
7use crate::type_object::PyTypeInfo;
8use crate::types::any::PyAnyMethods;
9use crate::types::{PyAny, PyDict, PyList, PyType};
10use crate::{ffi, Py, PyTypeCheck, Python};
11
12/// Represents a reference to a Python object supporting the mapping protocol.
13///
14/// Values of this type are accessed via PyO3's smart pointers, e.g. as
15/// [`Py<PyMapping>`][crate::Py] or [`Bound<'py, PyMapping>`][Bound].
16///
17/// For APIs available on mapping objects, see the [`PyMappingMethods`] trait which is implemented for
18/// [`Bound<'py, PyMapping>`][Bound].
19#[repr(transparent)]
20pub struct PyMapping(PyAny);
21pyobject_native_type_named!(PyMapping);
22
23impl PyMapping {
24    /// Register a pyclass as a subclass of `collections.abc.Mapping` (from the Python standard
25    /// library). This is equivalent to `collections.abc.Mapping.register(T)` in Python.
26    /// This registration is required for a pyclass to be castable from `PyAny` to `PyMapping`.
27    pub fn register<T: PyTypeInfo>(py: Python<'_>) -> PyResult<()> {
28        let ty = T::type_object(py);
29        get_mapping_abc(py)?.call_method1("register", (ty,))?;
30        Ok(())
31    }
32}
33
34/// Implementation of functionality for [`PyMapping`].
35///
36/// These methods are defined for the `Bound<'py, PyMapping>` smart pointer, so to use method call
37/// syntax these methods are separated into a trait, because stable Rust does not yet support
38/// `arbitrary_self_types`.
39#[doc(alias = "PyMapping")]
40pub trait PyMappingMethods<'py>: crate::sealed::Sealed {
41    /// Returns the number of objects in the mapping.
42    ///
43    /// This is equivalent to the Python expression `len(self)`.
44    fn len(&self) -> PyResult<usize>;
45
46    /// Returns whether the mapping is empty.
47    fn is_empty(&self) -> PyResult<bool>;
48
49    /// Determines if the mapping contains the specified key.
50    ///
51    /// This is equivalent to the Python expression `key in self`.
52    fn contains<K>(&self, key: K) -> PyResult<bool>
53    where
54        K: IntoPyObject<'py>;
55
56    /// Gets the item in self with key `key`.
57    ///
58    /// Returns an `Err` if the item with specified key is not found, usually `KeyError`.
59    ///
60    /// This is equivalent to the Python expression `self[key]`.
61    fn get_item<K>(&self, key: K) -> PyResult<Bound<'py, PyAny>>
62    where
63        K: IntoPyObject<'py>;
64
65    /// Sets the item in self with key `key`.
66    ///
67    /// This is equivalent to the Python expression `self[key] = value`.
68    fn set_item<K, V>(&self, key: K, value: V) -> PyResult<()>
69    where
70        K: IntoPyObject<'py>,
71        V: IntoPyObject<'py>;
72
73    /// Deletes the item with key `key`.
74    ///
75    /// This is equivalent to the Python statement `del self[key]`.
76    fn del_item<K>(&self, key: K) -> PyResult<()>
77    where
78        K: IntoPyObject<'py>;
79
80    /// Returns a list containing all keys in the mapping.
81    fn keys(&self) -> PyResult<Bound<'py, PyList>>;
82
83    /// Returns a list containing all values in the mapping.
84    fn values(&self) -> PyResult<Bound<'py, PyList>>;
85
86    /// Returns a list of all (key, value) pairs in the mapping.
87    fn items(&self) -> PyResult<Bound<'py, PyList>>;
88}
89
90impl<'py> PyMappingMethods<'py> for Bound<'py, PyMapping> {
91    #[inline]
92    fn len(&self) -> PyResult<usize> {
93        let v = unsafe { ffi::PyMapping_Size(self.as_ptr()) };
94        crate::err::error_on_minusone(self.py(), v)?;
95        Ok(v as usize)
96    }
97
98    #[inline]
99    fn is_empty(&self) -> PyResult<bool> {
100        self.len().map(|l| l == 0)
101    }
102
103    fn contains<K>(&self, key: K) -> PyResult<bool>
104    where
105        K: IntoPyObject<'py>,
106    {
107        PyAnyMethods::contains(&**self, key)
108    }
109
110    #[inline]
111    fn get_item<K>(&self, key: K) -> PyResult<Bound<'py, PyAny>>
112    where
113        K: IntoPyObject<'py>,
114    {
115        PyAnyMethods::get_item(&**self, key)
116    }
117
118    #[inline]
119    fn set_item<K, V>(&self, key: K, value: V) -> PyResult<()>
120    where
121        K: IntoPyObject<'py>,
122        V: IntoPyObject<'py>,
123    {
124        PyAnyMethods::set_item(&**self, key, value)
125    }
126
127    #[inline]
128    fn del_item<K>(&self, key: K) -> PyResult<()>
129    where
130        K: IntoPyObject<'py>,
131    {
132        PyAnyMethods::del_item(&**self, key)
133    }
134
135    #[inline]
136    fn keys(&self) -> PyResult<Bound<'py, PyList>> {
137        unsafe {
138            ffi::PyMapping_Keys(self.as_ptr())
139                .assume_owned_or_err(self.py())
140                .cast_into_unchecked()
141        }
142    }
143
144    #[inline]
145    fn values(&self) -> PyResult<Bound<'py, PyList>> {
146        unsafe {
147            ffi::PyMapping_Values(self.as_ptr())
148                .assume_owned_or_err(self.py())
149                .cast_into_unchecked()
150        }
151    }
152
153    #[inline]
154    fn items(&self) -> PyResult<Bound<'py, PyList>> {
155        unsafe {
156            ffi::PyMapping_Items(self.as_ptr())
157                .assume_owned_or_err(self.py())
158                .cast_into_unchecked()
159        }
160    }
161}
162
163fn get_mapping_abc(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> {
164    static MAPPING_ABC: PyOnceLock<Py<PyType>> = PyOnceLock::new();
165
166    MAPPING_ABC.import(py, "collections.abc", "Mapping")
167}
168
169impl PyTypeCheck for PyMapping {
170    const NAME: &'static str = "Mapping";
171    #[cfg(feature = "experimental-inspect")]
172    const PYTHON_TYPE: &'static str = "collections.abc.Mapping";
173
174    #[inline]
175    fn type_check(object: &Bound<'_, PyAny>) -> bool {
176        // Using `is_instance` for `collections.abc.Mapping` is slow, so provide
177        // optimized case dict as a well-known mapping
178        PyDict::is_type_of(object)
179            || get_mapping_abc(object.py())
180                .and_then(|abc| object.is_instance(abc))
181                .unwrap_or_else(|err| {
182                    err.write_unraisable(object.py(), Some(object));
183                    false
184                })
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use std::collections::HashMap;
191
192    use crate::{exceptions::PyKeyError, types::PyTuple};
193
194    use super::*;
195    use crate::conversion::IntoPyObject;
196
197    #[test]
198    fn test_len() {
199        Python::attach(|py| {
200            let mut v = HashMap::<i32, i32>::new();
201            let ob = (&v).into_pyobject(py).unwrap();
202            let mapping = ob.cast::<PyMapping>().unwrap();
203            assert_eq!(0, mapping.len().unwrap());
204            assert!(mapping.is_empty().unwrap());
205
206            v.insert(7, 32);
207            let ob = v.into_pyobject(py).unwrap();
208            let mapping2 = ob.cast::<PyMapping>().unwrap();
209            assert_eq!(1, mapping2.len().unwrap());
210            assert!(!mapping2.is_empty().unwrap());
211        });
212    }
213
214    #[test]
215    fn test_contains() {
216        Python::attach(|py| {
217            let mut v = HashMap::new();
218            v.insert("key0", 1234);
219            let ob = v.into_pyobject(py).unwrap();
220            let mapping = ob.cast::<PyMapping>().unwrap();
221            mapping.set_item("key1", "foo").unwrap();
222
223            assert!(mapping.contains("key0").unwrap());
224            assert!(mapping.contains("key1").unwrap());
225            assert!(!mapping.contains("key2").unwrap());
226        });
227    }
228
229    #[test]
230    fn test_get_item() {
231        Python::attach(|py| {
232            let mut v = HashMap::new();
233            v.insert(7, 32);
234            let ob = v.into_pyobject(py).unwrap();
235            let mapping = ob.cast::<PyMapping>().unwrap();
236            assert_eq!(
237                32,
238                mapping.get_item(7i32).unwrap().extract::<i32>().unwrap()
239            );
240            assert!(mapping
241                .get_item(8i32)
242                .unwrap_err()
243                .is_instance_of::<PyKeyError>(py));
244        });
245    }
246
247    #[test]
248    fn test_set_item() {
249        Python::attach(|py| {
250            let mut v = HashMap::new();
251            v.insert(7, 32);
252            let ob = v.into_pyobject(py).unwrap();
253            let mapping = ob.cast::<PyMapping>().unwrap();
254            assert!(mapping.set_item(7i32, 42i32).is_ok()); // change
255            assert!(mapping.set_item(8i32, 123i32).is_ok()); // insert
256            assert_eq!(
257                42i32,
258                mapping.get_item(7i32).unwrap().extract::<i32>().unwrap()
259            );
260            assert_eq!(
261                123i32,
262                mapping.get_item(8i32).unwrap().extract::<i32>().unwrap()
263            );
264        });
265    }
266
267    #[test]
268    fn test_del_item() {
269        Python::attach(|py| {
270            let mut v = HashMap::new();
271            v.insert(7, 32);
272            let ob = v.into_pyobject(py).unwrap();
273            let mapping = ob.cast::<PyMapping>().unwrap();
274            assert!(mapping.del_item(7i32).is_ok());
275            assert_eq!(0, mapping.len().unwrap());
276            assert!(mapping
277                .get_item(7i32)
278                .unwrap_err()
279                .is_instance_of::<PyKeyError>(py));
280        });
281    }
282
283    #[test]
284    fn test_items() {
285        Python::attach(|py| {
286            let mut v = HashMap::new();
287            v.insert(7, 32);
288            v.insert(8, 42);
289            v.insert(9, 123);
290            let ob = v.into_pyobject(py).unwrap();
291            let mapping = ob.cast::<PyMapping>().unwrap();
292            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
293            let mut key_sum = 0;
294            let mut value_sum = 0;
295            for el in mapping.items().unwrap().try_iter().unwrap() {
296                let tuple = el.unwrap().cast_into::<PyTuple>().unwrap();
297                key_sum += tuple.get_item(0).unwrap().extract::<i32>().unwrap();
298                value_sum += tuple.get_item(1).unwrap().extract::<i32>().unwrap();
299            }
300            assert_eq!(7 + 8 + 9, key_sum);
301            assert_eq!(32 + 42 + 123, value_sum);
302        });
303    }
304
305    #[test]
306    fn test_keys() {
307        Python::attach(|py| {
308            let mut v = HashMap::new();
309            v.insert(7, 32);
310            v.insert(8, 42);
311            v.insert(9, 123);
312            let ob = v.into_pyobject(py).unwrap();
313            let mapping = ob.cast::<PyMapping>().unwrap();
314            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
315            let mut key_sum = 0;
316            for el in mapping.keys().unwrap().try_iter().unwrap() {
317                key_sum += el.unwrap().extract::<i32>().unwrap();
318            }
319            assert_eq!(7 + 8 + 9, key_sum);
320        });
321    }
322
323    #[test]
324    fn test_values() {
325        Python::attach(|py| {
326            let mut v = HashMap::new();
327            v.insert(7, 32);
328            v.insert(8, 42);
329            v.insert(9, 123);
330            let ob = v.into_pyobject(py).unwrap();
331            let mapping = ob.cast::<PyMapping>().unwrap();
332            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
333            let mut values_sum = 0;
334            for el in mapping.values().unwrap().try_iter().unwrap() {
335                values_sum += el.unwrap().extract::<i32>().unwrap();
336            }
337            assert_eq!(32 + 42 + 123, values_sum);
338        });
339    }
340}