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::GILOnceCell;
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 downcastable 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                .downcast_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                .downcast_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                .downcast_into_unchecked()
159        }
160    }
161}
162
163fn get_mapping_abc(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> {
164    static MAPPING_ABC: GILOnceCell<Py<PyType>> = GILOnceCell::new();
165
166    MAPPING_ABC.import(py, "collections.abc", "Mapping")
167}
168
169impl PyTypeCheck for PyMapping {
170    const NAME: &'static str = "Mapping";
171
172    #[inline]
173    fn type_check(object: &Bound<'_, PyAny>) -> bool {
174        // Using `is_instance` for `collections.abc.Mapping` is slow, so provide
175        // optimized case dict as a well-known mapping
176        PyDict::is_type_of(object)
177            || get_mapping_abc(object.py())
178                .and_then(|abc| object.is_instance(abc))
179                .unwrap_or_else(|err| {
180                    err.write_unraisable(object.py(), Some(object));
181                    false
182                })
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use std::collections::HashMap;
189
190    use crate::{exceptions::PyKeyError, types::PyTuple};
191
192    use super::*;
193    use crate::conversion::IntoPyObject;
194
195    #[test]
196    fn test_len() {
197        Python::with_gil(|py| {
198            let mut v = HashMap::<i32, i32>::new();
199            let ob = (&v).into_pyobject(py).unwrap();
200            let mapping = ob.downcast::<PyMapping>().unwrap();
201            assert_eq!(0, mapping.len().unwrap());
202            assert!(mapping.is_empty().unwrap());
203
204            v.insert(7, 32);
205            let ob = v.into_pyobject(py).unwrap();
206            let mapping2 = ob.downcast::<PyMapping>().unwrap();
207            assert_eq!(1, mapping2.len().unwrap());
208            assert!(!mapping2.is_empty().unwrap());
209        });
210    }
211
212    #[test]
213    fn test_contains() {
214        Python::with_gil(|py| {
215            let mut v = HashMap::new();
216            v.insert("key0", 1234);
217            let ob = v.into_pyobject(py).unwrap();
218            let mapping = ob.downcast::<PyMapping>().unwrap();
219            mapping.set_item("key1", "foo").unwrap();
220
221            assert!(mapping.contains("key0").unwrap());
222            assert!(mapping.contains("key1").unwrap());
223            assert!(!mapping.contains("key2").unwrap());
224        });
225    }
226
227    #[test]
228    fn test_get_item() {
229        Python::with_gil(|py| {
230            let mut v = HashMap::new();
231            v.insert(7, 32);
232            let ob = v.into_pyobject(py).unwrap();
233            let mapping = ob.downcast::<PyMapping>().unwrap();
234            assert_eq!(
235                32,
236                mapping.get_item(7i32).unwrap().extract::<i32>().unwrap()
237            );
238            assert!(mapping
239                .get_item(8i32)
240                .unwrap_err()
241                .is_instance_of::<PyKeyError>(py));
242        });
243    }
244
245    #[test]
246    fn test_set_item() {
247        Python::with_gil(|py| {
248            let mut v = HashMap::new();
249            v.insert(7, 32);
250            let ob = v.into_pyobject(py).unwrap();
251            let mapping = ob.downcast::<PyMapping>().unwrap();
252            assert!(mapping.set_item(7i32, 42i32).is_ok()); // change
253            assert!(mapping.set_item(8i32, 123i32).is_ok()); // insert
254            assert_eq!(
255                42i32,
256                mapping.get_item(7i32).unwrap().extract::<i32>().unwrap()
257            );
258            assert_eq!(
259                123i32,
260                mapping.get_item(8i32).unwrap().extract::<i32>().unwrap()
261            );
262        });
263    }
264
265    #[test]
266    fn test_del_item() {
267        Python::with_gil(|py| {
268            let mut v = HashMap::new();
269            v.insert(7, 32);
270            let ob = v.into_pyobject(py).unwrap();
271            let mapping = ob.downcast::<PyMapping>().unwrap();
272            assert!(mapping.del_item(7i32).is_ok());
273            assert_eq!(0, mapping.len().unwrap());
274            assert!(mapping
275                .get_item(7i32)
276                .unwrap_err()
277                .is_instance_of::<PyKeyError>(py));
278        });
279    }
280
281    #[test]
282    fn test_items() {
283        Python::with_gil(|py| {
284            let mut v = HashMap::new();
285            v.insert(7, 32);
286            v.insert(8, 42);
287            v.insert(9, 123);
288            let ob = v.into_pyobject(py).unwrap();
289            let mapping = ob.downcast::<PyMapping>().unwrap();
290            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
291            let mut key_sum = 0;
292            let mut value_sum = 0;
293            for el in mapping.items().unwrap().try_iter().unwrap() {
294                let tuple = el.unwrap().downcast_into::<PyTuple>().unwrap();
295                key_sum += tuple.get_item(0).unwrap().extract::<i32>().unwrap();
296                value_sum += tuple.get_item(1).unwrap().extract::<i32>().unwrap();
297            }
298            assert_eq!(7 + 8 + 9, key_sum);
299            assert_eq!(32 + 42 + 123, value_sum);
300        });
301    }
302
303    #[test]
304    fn test_keys() {
305        Python::with_gil(|py| {
306            let mut v = HashMap::new();
307            v.insert(7, 32);
308            v.insert(8, 42);
309            v.insert(9, 123);
310            let ob = v.into_pyobject(py).unwrap();
311            let mapping = ob.downcast::<PyMapping>().unwrap();
312            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
313            let mut key_sum = 0;
314            for el in mapping.keys().unwrap().try_iter().unwrap() {
315                key_sum += el.unwrap().extract::<i32>().unwrap();
316            }
317            assert_eq!(7 + 8 + 9, key_sum);
318        });
319    }
320
321    #[test]
322    fn test_values() {
323        Python::with_gil(|py| {
324            let mut v = HashMap::new();
325            v.insert(7, 32);
326            v.insert(8, 42);
327            v.insert(9, 123);
328            let ob = v.into_pyobject(py).unwrap();
329            let mapping = ob.downcast::<PyMapping>().unwrap();
330            // Can't just compare against a vector of tuples since we don't have a guaranteed ordering.
331            let mut values_sum = 0;
332            for el in mapping.values().unwrap().try_iter().unwrap() {
333                values_sum += el.unwrap().extract::<i32>().unwrap();
334            }
335            assert_eq!(32 + 42 + 123, values_sum);
336        });
337    }
338}