1use crate::ffi_ptr_ext::FfiPtrExt;
3use crate::impl_::callback::IntoPyCallbackOutput;
4use crate::impl_::pyclass::{PyClassBaseType, PyClassDict, PyClassThreadChecker, PyClassWeakRef};
5use crate::impl_::pyclass_init::{PyNativeTypeInitializer, PyObjectInit};
6use crate::{ffi, Bound, Py, PyClass, PyResult, Python};
7use crate::{
8 ffi::PyTypeObject,
9 pycell::impl_::{PyClassBorrowChecker, PyClassMutability, PyClassObjectContents},
10};
11use std::{
12 cell::UnsafeCell,
13 marker::PhantomData,
14 mem::{ManuallyDrop, MaybeUninit},
15};
16
17pub struct PyClassInitializer<T: PyClass>(PyClassInitializerImpl<T>);
67
68enum PyClassInitializerImpl<T: PyClass> {
69 Existing(Py<T>),
70 New {
71 init: T,
72 super_init: <T::BaseType as PyClassBaseType>::Initializer,
73 },
74}
75
76impl<T: PyClass> PyClassInitializer<T> {
77 #[track_caller]
81 #[inline]
82 pub fn new(init: T, super_init: <T::BaseType as PyClassBaseType>::Initializer) -> Self {
83 assert!(
85 super_init.can_be_subclassed(),
86 "you cannot add a subclass to an existing value",
87 );
88 Self(PyClassInitializerImpl::New { init, super_init })
89 }
90
91 #[track_caller]
138 #[inline]
139 pub fn add_subclass<S>(self, subclass_value: S) -> PyClassInitializer<S>
140 where
141 S: PyClass<BaseType = T>,
142 S::BaseType: PyClassBaseType<Initializer = Self>,
143 {
144 PyClassInitializer::new(subclass_value, self)
145 }
146
147 pub(crate) fn create_class_object(self, py: Python<'_>) -> PyResult<Bound<'_, T>>
149 where
150 T: PyClass,
151 {
152 unsafe { self.create_class_object_of_type(py, T::type_object_raw(py)) }
153 }
154
155 pub(crate) unsafe fn create_class_object_of_type(
160 self,
161 py: Python<'_>,
162 target_type: *mut crate::ffi::PyTypeObject,
163 ) -> PyResult<Bound<'_, T>>
164 where
165 T: PyClass,
166 {
167 #[repr(C)]
170 struct PartiallyInitializedClassObject<T: PyClass> {
171 _ob_base: <T::BaseType as PyClassBaseType>::LayoutAsBase,
172 contents: MaybeUninit<PyClassObjectContents<T>>,
173 }
174
175 let (init, super_init) = match self.0 {
176 PyClassInitializerImpl::Existing(value) => return Ok(value.into_bound(py)),
177 PyClassInitializerImpl::New { init, super_init } => (init, super_init),
178 };
179
180 let obj = unsafe { super_init.into_new_object(py, target_type)? };
181
182 let part_init: *mut PartiallyInitializedClassObject<T> = obj.cast();
183 unsafe {
184 std::ptr::write(
185 (*part_init).contents.as_mut_ptr(),
186 PyClassObjectContents {
187 value: ManuallyDrop::new(UnsafeCell::new(init)),
188 borrow_checker: <T::PyClassMutability as PyClassMutability>::Storage::new(),
189 thread_checker: T::ThreadChecker::new(),
190 dict: T::Dict::INIT,
191 weakref: T::WeakRef::INIT,
192 },
193 );
194 }
195
196 Ok(unsafe { obj.assume_owned(py).cast_into_unchecked() })
199 }
200}
201
202impl<T: PyClass> PyObjectInit<T> for PyClassInitializer<T> {
203 unsafe fn into_new_object(
204 self,
205 py: Python<'_>,
206 subtype: *mut PyTypeObject,
207 ) -> PyResult<*mut ffi::PyObject> {
208 unsafe {
209 self.create_class_object_of_type(py, subtype)
210 .map(Bound::into_ptr)
211 }
212 }
213
214 #[inline]
215 fn can_be_subclassed(&self) -> bool {
216 !matches!(self.0, PyClassInitializerImpl::Existing(..))
217 }
218}
219
220impl<T> From<T> for PyClassInitializer<T>
221where
222 T: PyClass,
223 T::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<T::BaseType>>,
224{
225 #[inline]
226 fn from(value: T) -> PyClassInitializer<T> {
227 Self::new(value, PyNativeTypeInitializer(PhantomData))
228 }
229}
230
231impl<S, B> From<(S, B)> for PyClassInitializer<S>
232where
233 S: PyClass<BaseType = B>,
234 B: PyClass + PyClassBaseType<Initializer = PyClassInitializer<B>>,
235 B::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<B::BaseType>>,
236{
237 #[track_caller]
238 #[inline]
239 fn from(sub_and_base: (S, B)) -> PyClassInitializer<S> {
240 let (sub, base) = sub_and_base;
241 PyClassInitializer::from(base).add_subclass(sub)
242 }
243}
244
245impl<T: PyClass> From<Py<T>> for PyClassInitializer<T> {
246 #[inline]
247 fn from(value: Py<T>) -> PyClassInitializer<T> {
248 PyClassInitializer(PyClassInitializerImpl::Existing(value))
249 }
250}
251
252impl<'py, T: PyClass> From<Bound<'py, T>> for PyClassInitializer<T> {
253 #[inline]
254 fn from(value: Bound<'py, T>) -> PyClassInitializer<T> {
255 PyClassInitializer::from(value.unbind())
256 }
257}
258
259impl<T, U> IntoPyCallbackOutput<'_, PyClassInitializer<T>> for U
262where
263 T: PyClass,
264 U: Into<PyClassInitializer<T>>,
265{
266 #[inline]
267 fn convert(self, _py: Python<'_>) -> PyResult<PyClassInitializer<T>> {
268 Ok(self.into())
269 }
270}
271
272#[cfg(all(test, feature = "macros"))]
273mod tests {
274 use crate::prelude::*;
277
278 #[pyclass(crate = "crate", subclass)]
279 struct BaseClass {}
280
281 #[pyclass(crate = "crate", extends=BaseClass)]
282 struct SubClass {
283 _data: i32,
284 }
285
286 #[test]
287 #[should_panic]
288 fn add_subclass_to_py_is_unsound() {
289 Python::attach(|py| {
290 let base = Py::new(py, BaseClass {}).unwrap();
291 let _subclass = PyClassInitializer::from(base).add_subclass(SubClass { _data: 42 });
292 });
293 }
294}