1use 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
19pub trait PyObjectInit<T>: Sized {
24 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
38pub 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 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 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
105pub 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 #[track_caller]
169 #[inline]
170 pub fn new(init: T, super_init: <T::BaseType as PyClassBaseType>::Initializer) -> Self {
171 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 #[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 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 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 #[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 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
345impl<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 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}