pyo3/src/class/typeob.rs

// Copyright (c) 2017-present PyO3 Project and Contributors

use std::mem;
use std::ffi::CString;
use std::collections::HashMap;

use ::{ffi, class, PyErr, Python, PyResult, PythonObject};
use objects::PyType;
use callback::AbortOnDrop;
use class::{BaseObject, PyMethodDefType};


/// Trait implemented by Python object types that have a corresponding type object.
pub trait PyTypeObject {

    /// Retrieves the type object for this Python object type.
    fn type_object(py: Python) -> PyType;

}

/// Trait implemented by object that generated by py::class macro
pub trait PyTypeObjectInfo {

    fn type_name() -> &'static str;

    fn type_object() -> &'static mut ffi::PyTypeObject;

}


impl<T> PyTypeObject for T where T: BaseObject + PythonObject + PyTypeObjectInfo {

    #[inline]
    fn type_object(py: Python) -> PyType {
        let mut ty = <T as PyTypeObjectInfo>::type_object();

        if (ty.tp_flags & ffi::Py_TPFLAGS_READY) != 0 {
            unsafe { PyType::from_type_ptr(py, ty) }
        } else {
            // automatically initialize the class on-demand
            initialize_type::<T>(py, None, <T as PyTypeObjectInfo>::type_name(), ty).expect(
                format!("An error occurred while initializing class {}",
                        <T as PyTypeObjectInfo>::type_name()).as_ref());
            unsafe { PyType::from_type_ptr(py, ty) }
        }
    }
}

pub fn initialize_type<T>(py: Python, module_name: Option<&str>, type_name: &str,
                          type_object: &mut ffi::PyTypeObject) -> PyResult<PyType>
    where T: BaseObject + PythonObject
{
    // type name
    let name = match module_name {
        Some(module_name) => CString::new(format!("{}.{}", module_name, type_name)),
        None => CString::new(type_name)
    };
    let name = name.expect(
        "Module name/type name must not contain NUL byte").into_raw();

    type_object.tp_name = name;

    // dealloc
    type_object.tp_dealloc = Some(tp_dealloc_callback::<T>);

    // GC support
    <T as class::gc::PyGCProtocolImpl>::update_type_object(type_object);

    // type size
    type_object.tp_basicsize = <T as BaseObject>::size() as ffi::Py_ssize_t;

    // descriptor protocol
    type_object.tp_descr_get = class::descr::get_descrfunc::<T>();
    type_object.tp_descr_set = class::descr::set_descrfunc::<T>();

    // iterator methods
    <T as class::iter::PyIterProtocolImpl>::tp_as_iter(type_object);

    // number methods
    if let Some(meth) = ffi::PyNumberMethods::new::<T>() {
        static mut NB_METHODS: ffi::PyNumberMethods = ffi::PyNumberMethods_INIT;
        *(unsafe { &mut NB_METHODS }) = meth;
        type_object.tp_as_number = unsafe { &mut NB_METHODS };
        mem::forget(meth);
    } else {
        type_object.tp_as_number = 0 as *mut ffi::PyNumberMethods;
    }

    // mapping methods
    if let Some(meth) = <T as class::mapping::PyMappingProtocolImpl>::tp_as_mapping() {
        static mut MP_METHODS: ffi::PyMappingMethods = ffi::PyMappingMethods_INIT;
        *(unsafe { &mut MP_METHODS }) = meth;
        type_object.tp_as_mapping = unsafe { &mut MP_METHODS };
        mem::forget(meth);
    } else {
        type_object.tp_as_mapping = 0 as *mut ffi::PyMappingMethods;
    }

    // sequence methods
    if let Some(meth) = ffi::PySequenceMethods::new::<T>() {
        static mut SQ_METHODS: ffi::PySequenceMethods = ffi::PySequenceMethods_INIT;
        *(unsafe { &mut SQ_METHODS }) = meth;
        type_object.tp_as_sequence = unsafe { &mut SQ_METHODS };
        mem::forget(meth);
    } else {
        type_object.tp_as_sequence = 0 as *mut ffi::PySequenceMethods;
    }

    // async methods
    if let Some(meth) = <T as class::async::PyAsyncProtocolImpl>::tp_as_async() {
        static mut ASYNC_METHODS: ffi::PyAsyncMethods = ffi::PyAsyncMethods_INIT;
        *(unsafe { &mut ASYNC_METHODS }) = meth;
        type_object.tp_as_async = unsafe { &mut ASYNC_METHODS };
        mem::forget(meth);
    } else {
        type_object.tp_as_async = 0 as *mut ffi::PyAsyncMethods;
    }

    // buffer protocol
    if let Some(meth) = ffi::PyBufferProcs::new::<T>() {
        static mut BUFFER_PROCS: ffi::PyBufferProcs = ffi::PyBufferProcs_INIT;
        *(unsafe { &mut BUFFER_PROCS }) = meth;
        type_object.tp_as_buffer = unsafe { &mut BUFFER_PROCS };
        mem::forget(meth);
    } else {
        type_object.tp_as_buffer = 0 as *mut ffi::PyBufferProcs;
    }

    // normal methods
    let (new, call, mut methods) = py_class_method_defs::<T>();
    if !methods.is_empty() {
        methods.push(ffi::PyMethodDef_INIT);
        type_object.tp_methods = methods.as_ptr() as *mut _;

        static mut METHODS: *const ffi::PyMethodDef = 0 as *const _;
        *(unsafe { &mut METHODS }) = methods.as_ptr();

        mem::forget(methods);
    }
    // __new__ method
    type_object.tp_new = new;
    // __call__ method
    type_object.tp_call = call;

    // properties
    let mut props = py_class_properties::<T>();
    if !props.is_empty() {
        props.push(ffi::PyGetSetDef_INIT);
        let props = props.into_boxed_slice();
        type_object.tp_getset = props.as_ptr() as *mut _;

        static mut PROPS: *const ffi::PyGetSetDef = 0 as *const _;
        *(unsafe { &mut PROPS }) = props.as_ptr();

        // strange
        mem::forget(props);
    }

    // register type object
    unsafe {
        if ffi::PyType_Ready(type_object) == 0 {
            Ok(PyType::from_type_ptr(py, type_object))
        } else {
            Err(PyErr::fetch(py))
        }
    }
}

unsafe extern "C" fn tp_dealloc_callback<T>(obj: *mut ffi::PyObject) where T: BaseObject
{
    let guard = AbortOnDrop("Cannot unwind out of tp_dealloc");
    let py = Python::assume_gil_acquired();
    let r = T::dealloc(py, obj);
    mem::forget(guard);
    r
}

fn py_class_method_defs<T>() -> (Option<ffi::newfunc>,
                                 Option<ffi::PyCFunctionWithKeywords>,
                                 Vec<ffi::PyMethodDef>)  {
    let mut defs = Vec::new();
    let mut call = None;
    let mut new = None;

    for def in <T as class::number::PyNumberProtocolImpl>::py_methods() {
        match def {
            &PyMethodDefType::Method(ref def) => defs.push(def.as_method_def()),
            _ => (),
        }
    }
    for def in <T as class::methods::PyMethodsProtocolImpl>::py_methods() {
        match def {
            &PyMethodDefType::New(ref def) => {
                if let class::methods::PyMethodType::PyNewFunc(meth) = def.ml_meth {
                    new = Some(meth)
                }
            },
            &PyMethodDefType::Call(ref def) => {
                if let class::methods::PyMethodType::PyCFunctionWithKeywords(meth) = def.ml_meth {
                    call = Some(meth)
                } else {
                    panic!("Method type is not supoorted by tp_call slot")
                }
            }
            &PyMethodDefType::Method(ref def) => {
                defs.push(def.as_method_def())
            }
            _ => (),
        }
    }
    for def in <T as class::async::PyAsyncProtocolImpl>::methods() {
        defs.push(def.as_method_def())
    }
    for def in <T as class::context::PyContextProtocolImpl>::methods() {
        defs.push(def.as_method_def())
    }

    (new, call, defs)
}


fn py_class_properties<T>() -> Vec<ffi::PyGetSetDef> {
    let mut defs = HashMap::new();

    for def in <T as class::methods::PyMethodsProtocolImpl>::py_methods() {
        match def {
            &PyMethodDefType::Getter(ref getter) => {
                let name = getter.name.to_string();
                if !defs.contains_key(&name) {
                    let _ = defs.insert(name.clone(), ffi::PyGetSetDef_INIT);
                }
                let def = defs.get_mut(&name).unwrap();
                getter.copy_to(def);
            },
            &PyMethodDefType::Setter(ref setter) => {
                let name = setter.name.to_string();
                if !defs.contains_key(&name) {
                    let _ = defs.insert(name.clone(), ffi::PyGetSetDef_INIT);
                }
                let def = defs.get_mut(&name).unwrap();
                setter.copy_to(def);
            },
            _ => (),
        }
    }

    defs.values().map(|i| i.clone()).collect()
}