Skip to main content

rustpython_vm/builtins/
module.rs

1use super::{PyDict, PyDictRef, PyStr, PyStrRef, PyType, PyTypeRef};
2use crate::{
3    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
4    builtins::{PyStrInterned, pystr::AsPyStr},
5    class::PyClassImpl,
6    convert::ToPyObject,
7    function::{FuncArgs, PyMethodDef, PySetterValue},
8    import::{get_spec_file_origin, is_possibly_shadowing_path, is_stdlib_module_name},
9    types::{GetAttr, Initializer, Representable},
10};
11
12#[pyclass(module = false, name = "module")]
13#[derive(Debug)]
14pub struct PyModuleDef {
15    // pub index: usize,
16    pub name: &'static PyStrInterned,
17    pub doc: Option<&'static PyStrInterned>,
18    // pub size: isize,
19    pub methods: &'static [PyMethodDef],
20    pub slots: PyModuleSlots,
21    // traverse: traverse_proc
22    // clear: inquiry
23    // free: free_func
24}
25
26pub type ModuleCreate =
27    fn(&VirtualMachine, &PyObject, &'static PyModuleDef) -> PyResult<PyRef<PyModule>>;
28pub type ModuleExec = fn(&VirtualMachine, &Py<PyModule>) -> PyResult<()>;
29
30#[derive(Default)]
31pub struct PyModuleSlots {
32    pub create: Option<ModuleCreate>,
33    pub exec: Option<ModuleExec>,
34}
35
36impl core::fmt::Debug for PyModuleSlots {
37    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
38        f.debug_struct("PyModuleSlots")
39            .field("create", &self.create.is_some())
40            .field("exec", &self.exec.is_some())
41            .finish()
42    }
43}
44
45impl PyModuleDef {
46    /// Create a module from this definition (Phase 1 of multi-phase init).
47    ///
48    /// This performs:
49    /// 1. Create module object (using create slot if provided)
50    /// 2. Initialize module dict from def
51    /// 3. Add methods to module
52    ///
53    /// Does NOT add to sys.modules or call exec slot.
54    pub fn create_module(&'static self, vm: &VirtualMachine) -> PyResult<PyRef<PyModule>> {
55        use crate::PyPayload;
56
57        // Create module (use create slot if provided, else default creation)
58        let module = if let Some(create) = self.slots.create {
59            // Custom module creation
60            let spec = vm.ctx.new_str(self.name.as_str());
61            create(vm, spec.as_object(), self)?
62        } else {
63            // Default module creation
64            PyModule::from_def(self).into_ref(&vm.ctx)
65        };
66
67        // Initialize module dict and methods
68        PyModule::__init_dict_from_def(vm, &module);
69        module.__init_methods(vm)?;
70
71        Ok(module)
72    }
73
74    /// Execute the module's exec slot (Phase 2 of multi-phase init).
75    ///
76    /// Calls the exec slot if present. Returns Ok(()) if no exec slot.
77    pub fn exec_module(&'static self, vm: &VirtualMachine, module: &Py<PyModule>) -> PyResult<()> {
78        if let Some(exec) = self.slots.exec {
79            exec(vm, module)?;
80        }
81        Ok(())
82    }
83}
84
85#[allow(
86    clippy::new_without_default,
87    reason = "avoid a misleading Default implementation"
88)]
89#[pyclass(module = false, name = "module")]
90#[derive(Debug)]
91pub struct PyModule {
92    // PyObject *md_dict;
93    pub def: Option<&'static PyModuleDef>,
94    // state: Any
95    // weaklist
96    // for logging purposes after md_dict is cleared
97    pub name: Option<&'static PyStrInterned>,
98}
99
100impl PyPayload for PyModule {
101    #[inline]
102    fn class(ctx: &Context) -> &'static Py<PyType> {
103        ctx.types.module_type
104    }
105}
106
107#[derive(FromArgs)]
108pub struct ModuleInitArgs {
109    name: PyStrRef,
110    #[pyarg(any, default)]
111    doc: Option<PyStrRef>,
112}
113
114impl PyModule {
115    #[allow(clippy::new_without_default)]
116    pub const fn new() -> Self {
117        Self {
118            def: None,
119            name: None,
120        }
121    }
122
123    pub const fn from_def(def: &'static PyModuleDef) -> Self {
124        Self {
125            def: Some(def),
126            name: Some(def.name),
127        }
128    }
129
130    pub fn __init_dict_from_def(vm: &VirtualMachine, module: &Py<Self>) {
131        let doc = module.def.unwrap().doc.map(|doc| doc.to_owned());
132        module.init_dict(module.name.unwrap(), doc, vm);
133    }
134}
135
136impl Py<PyModule> {
137    pub fn __init_methods(&self, vm: &VirtualMachine) -> PyResult<()> {
138        debug_assert!(self.def.is_some());
139        for method in self.def.unwrap().methods {
140            let func = method
141                .to_function()
142                .with_module(self.name.unwrap())
143                .into_ref(&vm.ctx);
144            vm.__module_set_attr(self, vm.ctx.intern_str(method.name), func)?;
145        }
146        Ok(())
147    }
148
149    fn getattr_inner(&self, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
150        if let Some(attr) = self.as_object().generic_getattr_opt(name, None, vm)? {
151            return Ok(attr);
152        }
153        if let Ok(getattr) = self.dict().get_item(identifier!(vm, __getattr__), vm) {
154            return getattr.call((name.to_owned(),), vm);
155        }
156        let dict = self.dict();
157
158        // Get the raw __name__ object (may be a str subclass)
159        let mod_name_obj = dict
160            .get_item_opt(identifier!(vm, __name__), vm)
161            .ok()
162            .flatten();
163        let mod_name_str = mod_name_obj.as_ref().and_then(|n| {
164            n.downcast_ref::<PyStr>()
165                .map(|s| s.to_string_lossy().into_owned())
166        });
167
168        // If __name__ is not set or not a string, use a simpler error message
169        let mod_display = match mod_name_str.as_deref() {
170            Some(s) => s,
171            None => {
172                return Err(vm.new_attribute_error(format!("module has no attribute '{name}'")));
173            }
174        };
175
176        let spec = dict
177            .get_item_opt(vm.ctx.intern_str("__spec__"), vm)
178            .ok()
179            .flatten()
180            .filter(|s| !vm.is_none(s));
181
182        let origin = get_spec_file_origin(&spec, vm);
183
184        let is_possibly_shadowing = origin
185            .as_ref()
186            .map(|o| is_possibly_shadowing_path(o, vm))
187            .unwrap_or(false);
188        // Use the ORIGINAL __name__ object for stdlib check (may raise TypeError
189        // if __name__ is an unhashable str subclass)
190        let is_possibly_shadowing_stdlib = if is_possibly_shadowing {
191            if let Some(ref mod_name) = mod_name_obj {
192                is_stdlib_module_name(mod_name, vm)?
193            } else {
194                false
195            }
196        } else {
197            false
198        };
199
200        if is_possibly_shadowing_stdlib {
201            let origin = origin.as_ref().unwrap();
202            Err(vm.new_attribute_error(format!(
203                "module '{mod_display}' has no attribute '{name}' \
204                 (consider renaming '{origin}' since it has the same \
205                 name as the standard library module named '{mod_display}' \
206                 and prevents importing that standard library module)"
207            )))
208        } else {
209            let is_initializing = PyModule::is_initializing(&dict, vm);
210            if is_initializing {
211                if is_possibly_shadowing {
212                    let origin = origin.as_ref().unwrap();
213                    Err(vm.new_attribute_error(format!(
214                        "module '{mod_display}' has no attribute '{name}' \
215                         (consider renaming '{origin}' if it has the same name \
216                         as a library you intended to import)"
217                    )))
218                } else if let Some(ref origin) = origin {
219                    Err(vm.new_attribute_error(format!(
220                        "partially initialized module '{mod_display}' from '{origin}' \
221                         has no attribute '{name}' \
222                         (most likely due to a circular import)"
223                    )))
224                } else {
225                    Err(vm.new_attribute_error(format!(
226                        "partially initialized module '{mod_display}' \
227                         has no attribute '{name}' \
228                         (most likely due to a circular import)"
229                    )))
230                }
231            } else {
232                // Check for uninitialized submodule
233                let submodule_initializing =
234                    is_uninitialized_submodule(mod_name_str.as_ref(), name, vm);
235                if submodule_initializing {
236                    Err(vm.new_attribute_error(format!(
237                        "cannot access submodule '{name}' of module '{mod_display}' \
238                         (most likely due to a circular import)"
239                    )))
240                } else {
241                    Err(vm.new_attribute_error(format!(
242                        "module '{mod_display}' has no attribute '{name}'"
243                    )))
244                }
245            }
246        }
247    }
248
249    // TODO: to be replaced by the commented-out dict method above once dictoffset land
250    pub fn dict(&self) -> PyDictRef {
251        self.as_object().dict().unwrap()
252    }
253
254    // TODO: should be on PyModule, not Py<PyModule>
255    pub(crate) fn init_dict(
256        &self,
257        name: &'static PyStrInterned,
258        doc: Option<PyStrRef>,
259        vm: &VirtualMachine,
260    ) {
261        let dict = self.dict();
262        dict.set_item(identifier!(vm, __name__), name.to_object(), vm)
263            .expect("Failed to set __name__ on module");
264        dict.set_item(identifier!(vm, __doc__), doc.to_pyobject(vm), vm)
265            .expect("Failed to set __doc__ on module");
266        dict.set_item("__package__", vm.ctx.none(), vm)
267            .expect("Failed to set __package__ on module");
268        dict.set_item("__loader__", vm.ctx.none(), vm)
269            .expect("Failed to set __loader__ on module");
270        dict.set_item("__spec__", vm.ctx.none(), vm)
271            .expect("Failed to set __spec__ on module");
272    }
273
274    pub fn get_attr<'a>(&self, attr_name: impl AsPyStr<'a>, vm: &VirtualMachine) -> PyResult {
275        let attr_name = attr_name.as_pystr(&vm.ctx);
276        self.getattr_inner(attr_name, vm)
277    }
278
279    pub fn set_attr<'a>(
280        &self,
281        attr_name: impl AsPyStr<'a>,
282        attr_value: impl Into<PyObjectRef>,
283        vm: &VirtualMachine,
284    ) -> PyResult<()> {
285        self.as_object().set_attr(attr_name, attr_value, vm)
286    }
287}
288
289#[pyclass(
290    with(GetAttr, Initializer, Representable),
291    flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
292)]
293impl PyModule {
294    #[pyslot]
295    fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
296        Self::new().into_ref_with_type(vm, cls).map(Into::into)
297    }
298
299    #[pymethod]
300    fn __dir__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
301        // First check if __dict__ attribute exists and is actually a dictionary
302        let dict_attr = zelf.as_object().get_attr(identifier!(vm, __dict__), vm)?;
303        let dict = dict_attr
304            .downcast::<PyDict>()
305            .map_err(|_| vm.new_type_error("<module>.__dict__ is not a dictionary"))?;
306        let attrs = dict.into_iter().map(|(k, _v)| k).collect();
307        Ok(attrs)
308    }
309
310    #[pygetset]
311    fn __annotate__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
312        let dict = zelf.dict();
313        // Get __annotate__ from dict; if not present, insert None and return it
314        // See: module_get_annotate()
315        if let Some(annotate) = dict.get_item_opt(identifier!(vm, __annotate__), vm)? {
316            Ok(annotate)
317        } else {
318            let none = vm.ctx.none();
319            dict.set_item(identifier!(vm, __annotate__), none.clone(), vm)?;
320            Ok(none)
321        }
322    }
323
324    #[pygetset(setter)]
325    fn set___annotate__(
326        zelf: &Py<Self>,
327        value: PySetterValue,
328        vm: &VirtualMachine,
329    ) -> PyResult<()> {
330        match value {
331            PySetterValue::Assign(value) => {
332                if !vm.is_none(&value) && !value.is_callable() {
333                    return Err(vm.new_type_error("__annotate__ must be callable or None"));
334                }
335                let dict = zelf.dict();
336                dict.set_item(identifier!(vm, __annotate__), value.clone(), vm)?;
337                // Clear __annotations__ if value is not None
338                if !vm.is_none(&value) {
339                    dict.del_item(identifier!(vm, __annotations__), vm).ok();
340                }
341                Ok(())
342            }
343            PySetterValue::Delete => Err(vm.new_type_error("cannot delete __annotate__ attribute")),
344        }
345    }
346
347    #[pygetset]
348    fn __annotations__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
349        let dict = zelf.dict();
350
351        // Check if __annotations__ is already in dict (explicitly set)
352        if let Some(annotations) = dict.get_item_opt(identifier!(vm, __annotations__), vm)? {
353            return Ok(annotations);
354        }
355
356        // Check if module is initializing
357        let is_initializing = Self::is_initializing(&dict, vm);
358
359        // PEP 649: Get __annotate__ and call it if callable
360        let annotations = if let Some(annotate) =
361            dict.get_item_opt(identifier!(vm, __annotate__), vm)?
362            && annotate.is_callable()
363        {
364            // Call __annotate__(1) where 1 is FORMAT_VALUE
365            let result = annotate.call((1i32,), vm)?;
366            if !result.class().is(vm.ctx.types.dict_type) {
367                return Err(vm.new_type_error(format!(
368                    "__annotate__ returned non-dict of type '{}'",
369                    result.class().name()
370                )));
371            }
372            result
373        } else {
374            vm.ctx.new_dict().into()
375        };
376
377        // Cache result unless module is initializing
378        if !is_initializing {
379            dict.set_item(identifier!(vm, __annotations__), annotations.clone(), vm)?;
380        }
381
382        Ok(annotations)
383    }
384
385    /// Check if module is initializing via __spec__._initializing
386    fn is_initializing(dict: &PyDictRef, vm: &VirtualMachine) -> bool {
387        if let Ok(Some(spec)) = dict.get_item_opt(vm.ctx.intern_str("__spec__"), vm)
388            && let Ok(initializing) = spec.get_attr(vm.ctx.intern_str("_initializing"), vm)
389        {
390            return initializing.try_to_bool(vm).unwrap_or(false);
391        }
392        false
393    }
394
395    #[pygetset(setter)]
396    fn set___annotations__(
397        zelf: &Py<Self>,
398        value: PySetterValue,
399        vm: &VirtualMachine,
400    ) -> PyResult<()> {
401        let dict = zelf.dict();
402        match value {
403            PySetterValue::Assign(value) => {
404                dict.set_item(identifier!(vm, __annotations__), value, vm)?;
405                // Clear __annotate__ from dict
406                dict.del_item(identifier!(vm, __annotate__), vm).ok();
407                Ok(())
408            }
409            PySetterValue::Delete => {
410                if dict.del_item(identifier!(vm, __annotations__), vm).is_err() {
411                    return Err(vm.new_attribute_error("__annotations__"));
412                }
413                // Also clear __annotate__
414                dict.del_item(identifier!(vm, __annotate__), vm).ok();
415                Ok(())
416            }
417        }
418    }
419}
420
421impl Initializer for PyModule {
422    type Args = ModuleInitArgs;
423
424    fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
425        debug_assert!(
426            zelf.class()
427                .slots
428                .flags
429                .has_feature(crate::types::PyTypeFlags::HAS_DICT)
430        );
431        zelf.init_dict(vm.ctx.intern_str(args.name.as_wtf8()), args.doc, vm);
432        Ok(())
433    }
434}
435
436impl GetAttr for PyModule {
437    fn getattro(zelf: &Py<Self>, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
438        zelf.getattr_inner(name, vm)
439    }
440}
441
442impl Representable for PyModule {
443    #[inline]
444    fn repr(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyStrRef> {
445        // Use cached importlib reference (like interp->importlib)
446        let module_repr = vm.importlib.get_attr("_module_repr", vm)?;
447        let repr = module_repr.call((zelf.to_owned(),), vm)?;
448        repr.downcast()
449            .map_err(|_| vm.new_type_error("_module_repr did not return a string"))
450    }
451
452    #[cold]
453    fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
454        unreachable!("use repr instead")
455    }
456}
457
458pub(crate) fn init(context: &'static Context) {
459    PyModule::extend_class(context, context.types.module_type);
460}
461
462/// Check if {module_name}.{name} is an uninitialized submodule in sys.modules.
463fn is_uninitialized_submodule(
464    module_name: Option<&String>,
465    name: &Py<PyStr>,
466    vm: &VirtualMachine,
467) -> bool {
468    let mod_name = match module_name {
469        Some(n) => n.as_str(),
470        None => return false,
471    };
472    let full_name = format!("{mod_name}.{name}");
473    let sys_modules = match vm.sys_module.get_attr("modules", vm).ok() {
474        Some(m) => m,
475        None => return false,
476    };
477    let sub_mod = match sys_modules.get_item(&full_name, vm).ok() {
478        Some(m) => m,
479        None => return false,
480    };
481    let spec = match sub_mod.get_attr("__spec__", vm).ok() {
482        Some(s) if !vm.is_none(&s) => s,
483        _ => return false,
484    };
485    spec.get_attr("_initializing", vm)
486        .ok()
487        .and_then(|v| v.try_to_bool(vm).ok())
488        .unwrap_or(false)
489}