1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100

use super::objdict::PyDictRef;
use super::objstr::{PyString, PyStringRef};
use super::objtype::PyClassRef;
use crate::function::OptionalOption;
use crate::pyobject::{
    ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
};
use crate::vm::VirtualMachine;

#[pyclass]
#[derive(Debug)]
pub struct PyModule {}
pub type PyModuleRef = PyRef<PyModule>;

impl PyValue for PyModule {
    const HAVE_DICT: bool = true;

    fn class(vm: &VirtualMachine) -> PyClassRef {
        vm.ctx.module_type()
    }
}

pub fn init_module_dict(
    vm: &VirtualMachine,
    module_dict: &PyDictRef,
    name: PyObjectRef,
    doc: PyObjectRef,
) {
    module_dict
        .set_item("__name__", name, vm)
        .expect("Failed to set __name__ on module");
    module_dict
        .set_item("__doc__", doc, vm)
        .expect("Failed to set __doc__ on module");
    module_dict
        .set_item("__package__", vm.get_none(), vm)
        .expect("Failed to set __package__ on module");
    module_dict
        .set_item("__loader__", vm.get_none(), vm)
        .expect("Failed to set __loader__ on module");
    module_dict
        .set_item("__spec__", vm.get_none(), vm)
        .expect("Failed to set __spec__ on module");
}

#[pyimpl(flags(BASETYPE))]
impl PyModuleRef {
    #[pyslot]
    fn tp_new(
        cls: PyClassRef,
        name: PyStringRef,
        doc: OptionalOption<PyStringRef>,
        vm: &VirtualMachine,
    ) -> PyResult<PyModuleRef> {
        let zelf = PyModule {}.into_ref_with_type(vm, cls)?;
        init_module_dict(
            vm,
            &zelf.as_object().dict.as_ref().unwrap().borrow(),
            name.into_object(),
            doc.flat_option()
                .map_or_else(|| vm.get_none(), PyRef::into_object),
        );
        Ok(zelf)
    }

    fn name(self, vm: &VirtualMachine) -> Option<String> {
        vm.generic_getattribute(
            self.as_object().clone(),
            PyString::from("__name__").into_ref(vm),
        )
        .unwrap_or(None)
        .and_then(|obj| obj.payload::<PyString>().map(|s| s.as_str().to_owned()))
    }

    #[pymethod(magic)]
    fn getattribute(self, name: PyStringRef, vm: &VirtualMachine) -> PyResult {
        vm.generic_getattribute(self.as_object().clone(), name.clone())?
            .ok_or_else(|| {
                let module_name = if let Some(name) = self.name(vm) {
                    format!(" '{}'", name)
                } else {
                    "".to_owned()
                };
                vm.new_attribute_error(
                    format!("module{} has no attribute '{}'", module_name, name,),
                )
            })
    }

    #[pymethod(magic)]
    fn repr(self, vm: &VirtualMachine) -> PyResult {
        let importlib = vm.import("_frozen_importlib", &[], 0)?;
        let module_repr = vm.get_attribute(importlib, "_module_repr")?;
        vm.invoke(&module_repr, vec![self.into_object()])
    }
}

pub(crate) fn init(context: &PyContext) {
    PyModuleRef::extend_class(&context, &context.types.module_type);
}