Skip to main content

rustpython_vm/builtins/
module.rs

1use super::{PyDict, PyDictRef, PyStr, PyStrRef, PyType, PyTypeRef, PyUtf8Str};
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(crate) type ModuleCreate =
27    fn(&VirtualMachine, &PyObject, &'static PyModuleDef) -> PyResult<PyRef<PyModule>>;
28pub(crate) 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#[pyclass(module = false, name = "module")]
86// The dict lives in the object extension, not the payload. The offset is
87// the dict cell at the front of that extension.
88#[pymember(name = "__dict__", offset = ::rustpython_vm::object::dict_member_offset())]
89#[derive(Debug)]
90pub struct PyModule {
91    // PyObject *md_dict;
92    pub def: Option<&'static PyModuleDef>,
93    // state: Any
94    // weaklist
95    // for logging purposes after md_dict is cleared
96    pub name: Option<&'static PyStrInterned>,
97}
98
99impl PyPayload for PyModule {
100    #[inline]
101    fn class(ctx: &Context) -> &'static Py<PyType> {
102        ctx.types.module_type
103    }
104}
105
106#[derive(FromArgs)]
107pub struct ModuleInitArgs {
108    name: PyStrRef,
109    #[pyarg(any, default)]
110    doc: Option<PyStrRef>,
111}
112
113impl PyModule {
114    #[expect(
115        clippy::new_without_default,
116        reason = "avoid a misleading Default implementation"
117    )]
118    #[must_use]
119    pub const fn new() -> Self {
120        Self {
121            def: None,
122            name: None,
123        }
124    }
125
126    #[must_use]
127    pub const fn from_def(def: &'static PyModuleDef) -> Self {
128        Self {
129            def: Some(def),
130            name: Some(def.name),
131        }
132    }
133
134    pub fn __init_dict_from_def(vm: &VirtualMachine, module: &Py<Self>) {
135        let doc = module.def.unwrap().doc.map(|doc| doc.to_owned());
136        module.init_dict(module.name.unwrap(), doc, vm);
137    }
138}
139
140impl Py<PyModule> {
141    pub fn __init_methods(&self, vm: &VirtualMachine) -> PyResult<()> {
142        debug_assert!(self.def.is_some());
143        for method in self.def.unwrap().methods {
144            let func = method
145                .to_function()
146                .with_module(self.name.unwrap())
147                .with_module_object(self.to_owned().into())
148                .into_ref(&vm.ctx);
149            vm.__module_set_attr(self, vm.ctx.intern_str(method.name), func)?;
150        }
151        Ok(())
152    }
153
154    fn getattr_inner(&self, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
155        if let Some(attr) = self.as_object().generic_getattr_opt(name, None, vm)? {
156            return Ok(attr);
157        }
158        if let Ok(getattr) = self.dict().get_item(identifier!(vm, __getattr__), vm) {
159            return getattr.call((name.to_owned(),), vm);
160        }
161        let dict = self.dict();
162
163        // Get the raw __name__ object (may be a str subclass)
164        let mod_name_obj = dict
165            .get_item_opt(identifier!(vm, __name__), vm)
166            .ok()
167            .flatten();
168        let mod_name = mod_name_obj
169            .as_ref()
170            .and_then(|n| n.downcast_ref::<PyUtf8Str>());
171
172        // If __name__ is not set or not a string, use a simpler error message
173        let Some(mod_display) = mod_name.map(|s| s.as_str()) else {
174            return Err(vm.new_attribute_error(format!("module has no attribute '{name}'")));
175        };
176
177        let spec = dict
178            .get_item_opt(vm.ctx.intern_str("__spec__"), vm)
179            .ok()
180            .flatten()
181            .filter(|s| !vm.is_none(s));
182
183        let origin = get_spec_file_origin(spec.as_deref(), vm);
184
185        let is_possibly_shadowing = origin
186            .as_ref()
187            .is_some_and(|o| is_possibly_shadowing_path(o, vm));
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 = is_uninitialized_submodule(mod_name, name, vm);
234                if submodule_initializing {
235                    Err(vm.new_attribute_error(format!(
236                        "cannot access submodule '{name}' of module '{mod_display}' \
237                         (most likely due to a circular import)"
238                    )))
239                } else {
240                    Err(vm.new_attribute_error(format!(
241                        "module '{mod_display}' has no attribute '{name}'"
242                    )))
243                }
244            }
245        }
246    }
247
248    // TODO: to be replaced by the commented-out dict method above once dictoffset land
249    pub fn dict(&self) -> PyDictRef {
250        self.as_object().dict().unwrap()
251    }
252
253    // TODO: should be on PyModule, not Py<PyModule>
254    pub(crate) fn init_dict(
255        &self,
256        name: &'static PyStrInterned,
257        doc: Option<PyStrRef>,
258        vm: &VirtualMachine,
259    ) {
260        let dict = self.dict();
261        dict.set_item(identifier!(vm, __name__), name.to_object(), vm)
262            .expect("Failed to set __name__ on module");
263        dict.set_item(identifier!(vm, __doc__), doc.to_pyobject(vm), vm)
264            .expect("Failed to set __doc__ on module");
265        dict.set_item("__package__", vm.ctx.none(), vm)
266            .expect("Failed to set __package__ on module");
267        dict.set_item("__loader__", vm.ctx.none(), vm)
268            .expect("Failed to set __loader__ on module");
269        dict.set_item("__spec__", vm.ctx.none(), vm)
270            .expect("Failed to set __spec__ on module");
271    }
272
273    pub fn get_attr<'a>(&self, attr_name: impl AsPyStr<'a>, vm: &VirtualMachine) -> PyResult {
274        let attr_name = attr_name.as_pystr(&vm.ctx);
275        self.getattr_inner(attr_name, vm)
276    }
277
278    pub fn set_attr<'a>(
279        &self,
280        attr_name: impl AsPyStr<'a>,
281        attr_value: impl Into<PyObjectRef>,
282        vm: &VirtualMachine,
283    ) -> PyResult<()> {
284        self.as_object().set_attr(attr_name, attr_value, vm)
285    }
286}
287
288#[pyclass(
289    with(GetAttr, Initializer, Representable),
290    flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
291)]
292impl PyModule {
293    #[pyslot]
294    fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
295        Self::new().into_ref_with_type(vm, cls).map(Into::into)
296    }
297
298    #[pymethod]
299    fn __dir__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
300        // First check if __dict__ attribute exists and is actually a dictionary
301        let dict_attr = zelf.as_object().get_attr(identifier!(vm, __dict__), vm)?;
302        let dict = dict_attr
303            .downcast::<PyDict>()
304            .map_err(|_| vm.new_type_error("<module>.__dict__ is not a dictionary"))?;
305        // PEP 562: honor a module-level __dir__ if one is defined
306        if let Some(dir_func) = dict.get_item_opt(identifier!(vm, __dir__), vm)? {
307            return dir_func.call((), vm)?.try_to_value(vm);
308        }
309        let attrs = dict.into_iter().map(|(k, _v)| k).collect();
310        Ok(attrs)
311    }
312
313    #[pygetset]
314    fn __annotate__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
315        let dict = zelf.dict();
316        // Get __annotate__ from dict; if not present, insert None and return it
317        // See: module_get_annotate()
318        if let Some(annotate) = dict.get_item_opt(identifier!(vm, __annotate__), vm)? {
319            Ok(annotate)
320        } else {
321            let none = vm.ctx.none();
322            dict.set_item(identifier!(vm, __annotate__), none.clone(), vm)?;
323            Ok(none)
324        }
325    }
326
327    #[pygetset(setter)]
328    fn set___annotate__(
329        zelf: &Py<Self>,
330        value: PySetterValue,
331        vm: &VirtualMachine,
332    ) -> PyResult<()> {
333        match value {
334            PySetterValue::Assign(value) => {
335                if !vm.is_none(&value) && !value.is_callable() {
336                    return Err(vm.new_type_error("__annotate__ must be callable or None"));
337                }
338                let dict = zelf.dict();
339                dict.set_item(identifier!(vm, __annotate__), value.clone(), vm)?;
340                // Clear __annotations__ if value is not None
341                if !vm.is_none(&value) {
342                    dict.del_item(identifier!(vm, __annotations__), vm).ok();
343                }
344                Ok(())
345            }
346            PySetterValue::Delete => Err(vm.new_type_error("cannot delete __annotate__ attribute")),
347        }
348    }
349
350    #[pygetset]
351    fn __annotations__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
352        let dict = zelf.dict();
353
354        // Check if __annotations__ is already in dict (explicitly set)
355        if let Some(annotations) = dict.get_item_opt(identifier!(vm, __annotations__), vm)? {
356            return Ok(annotations);
357        }
358
359        // Check if module is initializing
360        let is_initializing = Self::is_initializing(&dict, vm);
361
362        // PEP 649: Get __annotate__ and call it if callable
363        let annotations = if let Some(annotate) =
364            dict.get_item_opt(identifier!(vm, __annotate__), vm)?
365            && annotate.is_callable()
366        {
367            // Call __annotate__(1) where 1 is FORMAT_VALUE
368            let result = annotate.call((1i32,), vm)?;
369            if !result.class().is(vm.ctx.types.dict_type) {
370                return Err(vm.new_type_error(format!(
371                    "__annotate__ returned non-dict of type '{}'",
372                    result.class().name()
373                )));
374            }
375            result
376        } else {
377            vm.ctx.new_dict().into()
378        };
379
380        // Cache result unless module is initializing
381        if !is_initializing {
382            dict.set_item(identifier!(vm, __annotations__), annotations.clone(), vm)?;
383        }
384
385        Ok(annotations)
386    }
387
388    /// Check if module is initializing via __spec__._initializing
389    fn is_initializing(dict: &Py<PyDict>, vm: &VirtualMachine) -> bool {
390        if let Ok(Some(spec)) = dict.get_item_opt(vm.ctx.intern_str("__spec__"), vm)
391            && let Ok(initializing) = spec.get_attr(vm.ctx.intern_str("_initializing"), vm)
392        {
393            return initializing.try_to_bool(vm).unwrap_or(false);
394        }
395        false
396    }
397
398    #[pygetset(setter)]
399    fn set___annotations__(
400        zelf: &Py<Self>,
401        value: PySetterValue,
402        vm: &VirtualMachine,
403    ) -> PyResult<()> {
404        let dict = zelf.dict();
405        match value {
406            PySetterValue::Assign(value) => {
407                dict.set_item(identifier!(vm, __annotations__), value, vm)?;
408                // Clear __annotate__ from dict
409                dict.del_item(identifier!(vm, __annotate__), vm).ok();
410                Ok(())
411            }
412            PySetterValue::Delete => {
413                if dict.del_item(identifier!(vm, __annotations__), vm).is_err() {
414                    return Err(vm.new_attribute_error("__annotations__"));
415                }
416                // Also clear __annotate__
417                dict.del_item(identifier!(vm, __annotate__), vm).ok();
418                Ok(())
419            }
420        }
421    }
422}
423
424impl Initializer for PyModule {
425    type Args = ModuleInitArgs;
426
427    fn init(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
428        debug_assert!(
429            zelf.class()
430                .slots
431                .flags
432                .has_feature(crate::types::PyTypeFlags::HAS_DICT)
433        );
434        zelf.init_dict(vm.ctx.intern_str(args.name.as_wtf8()), args.doc, vm);
435        Ok(())
436    }
437}
438
439impl GetAttr for PyModule {
440    fn getattro(zelf: &Py<Self>, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
441        zelf.getattr_inner(name, vm)
442    }
443}
444
445impl Representable for PyModule {
446    #[inline]
447    fn repr(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyStrRef> {
448        // Use cached importlib reference (like interp->importlib)
449        let module_repr = vm.importlib.get_attr("_module_repr", vm)?;
450        let repr = module_repr.call((zelf.to_owned(),), vm)?;
451        repr.downcast()
452            .map_err(|_| vm.new_type_error("_module_repr did not return a string"))
453    }
454
455    #[cold]
456    fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
457        unreachable!("use repr instead")
458    }
459}
460
461pub(crate) fn init(context: &'static Context) {
462    PyModule::extend_class(context, context.types.module_type);
463}
464
465/// Check if {module_name}.{name} is an uninitialized submodule in sys.modules.
466fn is_uninitialized_submodule(
467    module_name: Option<&Py<PyUtf8Str>>,
468    name: &Py<PyStr>,
469    vm: &VirtualMachine,
470) -> bool {
471    let Some(mod_name) = module_name else {
472        return false;
473    };
474
475    let Ok(sys_modules) = vm.sys_module.get_attr("modules", vm) else {
476        return false;
477    };
478
479    let full_name = vm.ctx.new_utf8_str(format!("{}.{name}", mod_name.as_str()));
480    let Ok(sub_mod) = sys_modules.get_item(&*full_name, vm) else {
481        return false;
482    };
483
484    let spec = match sub_mod.get_attr("__spec__", vm) {
485        Ok(s) if !vm.is_none(&s) => s,
486        _ => return false,
487    };
488
489    spec.get_attr("_initializing", vm)
490        .and_then(|v| v.try_to_bool(vm))
491        .unwrap_or(false)
492}