pyo3/
pyclass_init.rs

1//! Contains initialization utilities for `#[pyclass]`.
2use crate::callback::IntoPyCallbackOutput;
3use crate::ffi_ptr_ext::FfiPtrExt;
4use crate::impl_::pyclass::{PyClassBaseType, PyClassDict, PyClassThreadChecker, PyClassWeakRef};
5use crate::internal::get_slot::TP_ALLOC;
6use crate::types::{PyAnyMethods, PyType};
7use crate::{ffi, Borrowed, Bound, Py, PyClass, PyErr, PyResult, Python};
8use crate::{
9    ffi::PyTypeObject,
10    pycell::impl_::{PyClassBorrowChecker, PyClassMutability, PyClassObjectContents},
11    type_object::PyTypeInfo,
12};
13use std::{
14    cell::UnsafeCell,
15    marker::PhantomData,
16    mem::{ManuallyDrop, MaybeUninit},
17};
18
19/// Initializer for Python types.
20///
21/// This trait is intended to use internally for distinguishing `#[pyclass]` and
22/// Python native types.
23pub trait PyObjectInit<T>: Sized {
24    /// # Safety
25    /// - `subtype` must be a valid pointer to a type object of T or a subclass.
26    unsafe fn into_new_object(
27        self,
28        py: Python<'_>,
29        subtype: *mut PyTypeObject,
30    ) -> PyResult<*mut ffi::PyObject>;
31
32    #[doc(hidden)]
33    fn can_be_subclassed(&self) -> bool;
34
35    private_decl! {}
36}
37
38/// Initializer for Python native types, like `PyDict`.
39pub struct PyNativeTypeInitializer<T: PyTypeInfo>(PhantomData<T>);
40
41impl<T: PyTypeInfo> PyObjectInit<T> for PyNativeTypeInitializer<T> {
42    unsafe fn into_new_object(
43        self,
44        py: Python<'_>,
45        subtype: *mut PyTypeObject,
46    ) -> PyResult<*mut ffi::PyObject> {
47        unsafe fn inner(
48            py: Python<'_>,
49            type_object: *mut PyTypeObject,
50            subtype: *mut PyTypeObject,
51        ) -> PyResult<*mut ffi::PyObject> {
52            // HACK (due to FIXME below): PyBaseObject_Type's tp_new isn't happy with NULL arguments
53            let is_base_object = type_object == std::ptr::addr_of_mut!(ffi::PyBaseObject_Type);
54            let subtype_borrowed: Borrowed<'_, '_, PyType> = subtype
55                .cast::<ffi::PyObject>()
56                .assume_borrowed_unchecked(py)
57                .downcast_unchecked();
58
59            if is_base_object {
60                let alloc = subtype_borrowed
61                    .get_slot(TP_ALLOC)
62                    .unwrap_or(ffi::PyType_GenericAlloc);
63
64                let obj = alloc(subtype, 0);
65                return if obj.is_null() {
66                    Err(PyErr::fetch(py))
67                } else {
68                    Ok(obj)
69                };
70            }
71
72            #[cfg(Py_LIMITED_API)]
73            unreachable!("subclassing native types is not possible with the `abi3` feature");
74
75            #[cfg(not(Py_LIMITED_API))]
76            {
77                match (*type_object).tp_new {
78                    // FIXME: Call __new__ with actual arguments
79                    Some(newfunc) => {
80                        let obj = newfunc(subtype, std::ptr::null_mut(), std::ptr::null_mut());
81                        if obj.is_null() {
82                            Err(PyErr::fetch(py))
83                        } else {
84                            Ok(obj)
85                        }
86                    }
87                    None => Err(crate::exceptions::PyTypeError::new_err(
88                        "base type without tp_new",
89                    )),
90                }
91            }
92        }
93        let type_object = T::type_object_raw(py);
94        inner(py, type_object, subtype)
95    }
96
97    #[inline]
98    fn can_be_subclassed(&self) -> bool {
99        true
100    }
101
102    private_impl! {}
103}
104
105/// Initializer for our `#[pyclass]` system.
106///
107/// You can use this type to initialize complicatedly nested `#[pyclass]`.
108///
109/// # Examples
110///
111/// ```
112/// # use pyo3::prelude::*;
113/// # use pyo3::py_run;
114/// #[pyclass(subclass)]
115/// struct BaseClass {
116///     #[pyo3(get)]
117///     basename: &'static str,
118/// }
119/// #[pyclass(extends=BaseClass, subclass)]
120/// struct SubClass {
121///     #[pyo3(get)]
122///     subname: &'static str,
123/// }
124/// #[pyclass(extends=SubClass)]
125/// struct SubSubClass {
126///     #[pyo3(get)]
127///     subsubname: &'static str,
128/// }
129///
130/// #[pymethods]
131/// impl SubSubClass {
132///     #[new]
133///     fn new() -> PyClassInitializer<Self> {
134///         PyClassInitializer::from(BaseClass { basename: "base" })
135///             .add_subclass(SubClass { subname: "sub" })
136///             .add_subclass(SubSubClass {
137///                 subsubname: "subsub",
138///             })
139///     }
140/// }
141/// Python::with_gil(|py| {
142///     let typeobj = py.get_type_bound::<SubSubClass>();
143///     let sub_sub_class = typeobj.call((), None).unwrap();
144///     py_run!(
145///         py,
146///         sub_sub_class,
147///         r#"
148///  assert sub_sub_class.basename == 'base'
149///  assert sub_sub_class.subname == 'sub'
150///  assert sub_sub_class.subsubname == 'subsub'"#
151///     );
152/// });
153/// ```
154pub struct PyClassInitializer<T: PyClass>(PyClassInitializerImpl<T>);
155
156enum PyClassInitializerImpl<T: PyClass> {
157    Existing(Py<T>),
158    New {
159        init: T,
160        super_init: <T::BaseType as PyClassBaseType>::Initializer,
161    },
162}
163
164impl<T: PyClass> PyClassInitializer<T> {
165    /// Constructs a new initializer from value `T` and base class' initializer.
166    ///
167    /// It is recommended to use `add_subclass` instead of this method for most usage.
168    #[track_caller]
169    #[inline]
170    pub fn new(init: T, super_init: <T::BaseType as PyClassBaseType>::Initializer) -> Self {
171        // This is unsound; see https://github.com/PyO3/pyo3/issues/4452.
172        assert!(
173            super_init.can_be_subclassed(),
174            "you cannot add a subclass to an existing value",
175        );
176        Self(PyClassInitializerImpl::New { init, super_init })
177    }
178
179    /// Constructs a new initializer from an initializer for the base class.
180    ///
181    /// # Examples
182    /// ```
183    /// use pyo3::prelude::*;
184    ///
185    /// #[pyclass(subclass)]
186    /// struct BaseClass {
187    ///     #[pyo3(get)]
188    ///     value: i32,
189    /// }
190    ///
191    /// impl BaseClass {
192    ///     fn new(value: i32) -> PyResult<Self> {
193    ///         Ok(Self { value })
194    ///     }
195    /// }
196    ///
197    /// #[pyclass(extends=BaseClass)]
198    /// struct SubClass {}
199    ///
200    /// #[pymethods]
201    /// impl SubClass {
202    ///     #[new]
203    ///     fn new(value: i32) -> PyResult<PyClassInitializer<Self>> {
204    ///         let base_init = PyClassInitializer::from(BaseClass::new(value)?);
205    ///         Ok(base_init.add_subclass(SubClass {}))
206    ///     }
207    /// }
208    ///
209    /// fn main() -> PyResult<()> {
210    ///     Python::with_gil(|py| {
211    ///         let m = PyModule::new_bound(py, "example")?;
212    ///         m.add_class::<SubClass>()?;
213    ///         m.add_class::<BaseClass>()?;
214    ///
215    ///         let instance = m.getattr("SubClass")?.call1((92,))?;
216    ///
217    ///         // `SubClass` does not have a `value` attribute, but `BaseClass` does.
218    ///         let n = instance.getattr("value")?.extract::<i32>()?;
219    ///         assert_eq!(n, 92);
220    ///
221    ///         Ok(())
222    ///     })
223    /// }
224    /// ```
225    #[track_caller]
226    #[inline]
227    pub fn add_subclass<S>(self, subclass_value: S) -> PyClassInitializer<S>
228    where
229        S: PyClass<BaseType = T>,
230        S::BaseType: PyClassBaseType<Initializer = Self>,
231    {
232        PyClassInitializer::new(subclass_value, self)
233    }
234
235    /// Creates a new PyCell and initializes it.
236    pub(crate) fn create_class_object(self, py: Python<'_>) -> PyResult<Bound<'_, T>>
237    where
238        T: PyClass,
239    {
240        unsafe { self.create_class_object_of_type(py, T::type_object_raw(py)) }
241    }
242
243    /// Creates a new class object and initializes it given a typeobject `subtype`.
244    ///
245    /// # Safety
246    /// `subtype` must be a valid pointer to the type object of T or a subclass.
247    pub(crate) unsafe fn create_class_object_of_type(
248        self,
249        py: Python<'_>,
250        target_type: *mut crate::ffi::PyTypeObject,
251    ) -> PyResult<Bound<'_, T>>
252    where
253        T: PyClass,
254    {
255        /// Layout of a PyClassObject after base new has been called, but the contents have not yet been
256        /// written.
257        #[repr(C)]
258        struct PartiallyInitializedClassObject<T: PyClass> {
259            _ob_base: <T::BaseType as PyClassBaseType>::LayoutAsBase,
260            contents: MaybeUninit<PyClassObjectContents<T>>,
261        }
262
263        let (init, super_init) = match self.0 {
264            PyClassInitializerImpl::Existing(value) => return Ok(value.into_bound(py)),
265            PyClassInitializerImpl::New { init, super_init } => (init, super_init),
266        };
267
268        let obj = super_init.into_new_object(py, target_type)?;
269
270        let part_init: *mut PartiallyInitializedClassObject<T> = obj.cast();
271        std::ptr::write(
272            (*part_init).contents.as_mut_ptr(),
273            PyClassObjectContents {
274                value: ManuallyDrop::new(UnsafeCell::new(init)),
275                borrow_checker: <T::PyClassMutability as PyClassMutability>::Storage::new(),
276                thread_checker: T::ThreadChecker::new(),
277                dict: T::Dict::INIT,
278                weakref: T::WeakRef::INIT,
279            },
280        );
281
282        // Safety: obj is a valid pointer to an object of type `target_type`, which` is a known
283        // subclass of `T`
284        Ok(obj.assume_owned(py).downcast_into_unchecked())
285    }
286}
287
288impl<T: PyClass> PyObjectInit<T> for PyClassInitializer<T> {
289    unsafe fn into_new_object(
290        self,
291        py: Python<'_>,
292        subtype: *mut PyTypeObject,
293    ) -> PyResult<*mut ffi::PyObject> {
294        self.create_class_object_of_type(py, subtype)
295            .map(Bound::into_ptr)
296    }
297
298    #[inline]
299    fn can_be_subclassed(&self) -> bool {
300        !matches!(self.0, PyClassInitializerImpl::Existing(..))
301    }
302
303    private_impl! {}
304}
305
306impl<T> From<T> for PyClassInitializer<T>
307where
308    T: PyClass,
309    T::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<T::BaseType>>,
310{
311    #[inline]
312    fn from(value: T) -> PyClassInitializer<T> {
313        Self::new(value, PyNativeTypeInitializer(PhantomData))
314    }
315}
316
317impl<S, B> From<(S, B)> for PyClassInitializer<S>
318where
319    S: PyClass<BaseType = B>,
320    B: PyClass,
321    B::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<B::BaseType>>,
322{
323    #[track_caller]
324    #[inline]
325    fn from(sub_and_base: (S, B)) -> PyClassInitializer<S> {
326        let (sub, base) = sub_and_base;
327        PyClassInitializer::from(base).add_subclass(sub)
328    }
329}
330
331impl<T: PyClass> From<Py<T>> for PyClassInitializer<T> {
332    #[inline]
333    fn from(value: Py<T>) -> PyClassInitializer<T> {
334        PyClassInitializer(PyClassInitializerImpl::Existing(value))
335    }
336}
337
338impl<'py, T: PyClass> From<Bound<'py, T>> for PyClassInitializer<T> {
339    #[inline]
340    fn from(value: Bound<'py, T>) -> PyClassInitializer<T> {
341        PyClassInitializer::from(value.unbind())
342    }
343}
344
345// Implementation used by proc macros to allow anything convertible to PyClassInitializer<T> to be
346// the return value of pyclass #[new] method (optionally wrapped in `Result<U, E>`).
347impl<T, U> IntoPyCallbackOutput<PyClassInitializer<T>> for U
348where
349    T: PyClass,
350    U: Into<PyClassInitializer<T>>,
351{
352    #[inline]
353    fn convert(self, _py: Python<'_>) -> PyResult<PyClassInitializer<T>> {
354        Ok(self.into())
355    }
356}
357
358#[cfg(all(test, feature = "macros"))]
359mod tests {
360    //! See https://github.com/PyO3/pyo3/issues/4452.
361
362    use crate::prelude::*;
363
364    #[pyclass(crate = "crate", subclass)]
365    struct BaseClass {}
366
367    #[pyclass(crate = "crate", extends=BaseClass)]
368    struct SubClass {
369        _data: i32,
370    }
371
372    #[test]
373    #[should_panic]
374    fn add_subclass_to_py_is_unsound() {
375        Python::with_gil(|py| {
376            let base = Py::new(py, BaseClass {}).unwrap();
377            let _subclass = PyClassInitializer::from(base).add_subclass(SubClass { _data: 42 });
378        });
379    }
380}