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