Add sys._getframemodulename (#5873)

youknowone · web-flow · commit d5d350792103 · 2025-06-30T22:50:27.000+09:00
diff --git a/.cspell.dict/python-more.txt b/.cspell.dict/python-more.txt
@@ -78,6 +78,7 @@ getfilesystemencodeerrors
 getfilesystemencoding
 getformat
 getframe
+getframemodulename
 getnewargs
 getpip
 getrandom
diff --git a/extra_tests/snippets/stdlib_sys.py b/extra_tests/snippets/stdlib_sys.py
@@ -147,3 +147,14 @@ def recursive_call(n):
 proc = subprocess.run(args, stdout=subprocess.PIPE, universal_newlines=True, env=env)
 assert proc.stdout.rstrip() == "True"
 assert proc.returncode == 0, proc
+
+assert sys._getframemodulename() == "__main__", sys._getframemodulename()
+
+
+def test_getframemodulename():
+    return sys._getframemodulename()
+
+
+test_getframemodulename.__module__ = "awesome_module"
+
+assert test_getframemodulename() == "awesome_module"
diff --git a/vm/src/builtins/function.rs b/vm/src/builtins/function.rs
@@ -315,7 +315,9 @@ impl PyFunction {
 
         Ok(())
     }
+}
 
+impl Py<PyFunction> {
     pub fn invoke_with_locals(
         &self,
         func_args: FuncArgs,
@@ -352,6 +354,7 @@ impl PyFunction {
             Scope::new(Some(locals), self.globals.clone()),
             vm.builtins.dict(),
             self.closure.as_ref().map_or(&[], |c| c.as_slice()),
+            Some(self.to_owned().into()),
             vm,
         )
         .into_ref(&vm.ctx);
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -97,6 +97,7 @@ type Lasti = std::cell::Cell<u32>;
 #[pyclass(module = false, name = "frame")]
 pub struct Frame {
     pub code: PyRef<PyCode>,
+    pub func_obj: Option<PyObjectRef>,
 
     pub fastlocals: PyMutex<Box<[Option<PyObjectRef>]>>,
     pub(crate) cells_frees: Box<[PyCellRef]>,
@@ -139,6 +140,7 @@ impl Frame {
         scope: Scope,
         builtins: PyDictRef,
         closure: &[PyCellRef],
+        func_obj: Option<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> Frame {
         let cells_frees = std::iter::repeat_with(|| PyCell::default().into_ref(&vm.ctx))
@@ -160,6 +162,7 @@ impl Frame {
             globals: scope.globals,
             builtins,
             code,
+            func_obj,
             lasti: Lasti::new(0),
             state: PyMutex::new(state),
             trace: PyMutex::new(vm.ctx.none()),
diff --git a/vm/src/stdlib/sys.rs b/vm/src/stdlib/sys.rs
@@ -513,6 +513,32 @@ mod sys {
         Ok(frame.clone())
     }
 
+    #[pyfunction]
+    fn _getframemodulename(depth: OptionalArg<usize>, vm: &VirtualMachine) -> PyResult {
+        let depth = depth.into_option().unwrap_or(0);
+
+        // Get the frame at the specified depth
+        if depth > vm.frames.borrow().len() - 1 {
+            return Ok(vm.ctx.none());
+        }
+
+        let idx = vm.frames.borrow().len() - depth - 1;
+        let frame = &vm.frames.borrow()[idx];
+
+        // If the frame has a function object, return its __module__ attribute
+        if let Some(func_obj) = &frame.func_obj {
+            match func_obj.get_attr(identifier!(vm, __module__), vm) {
+                Ok(module) => Ok(module),
+                Err(_) => {
+                    // CPython clears the error and returns None
+                    Ok(vm.ctx.none())
+                }
+            }
+        } else {
+            Ok(vm.ctx.none())
+        }
+    }
+
     #[pyfunction]
     fn gettrace(vm: &VirtualMachine) -> PyObjectRef {
         vm.trace_func.borrow().clone()
diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -449,7 +449,26 @@ impl VirtualMachine {
     }
 
     pub fn run_code_obj(&self, code: PyRef<PyCode>, scope: Scope) -> PyResult {
-        let frame = Frame::new(code, scope, self.builtins.dict(), &[], self).into_ref(&self.ctx);
+        use crate::builtins::PyFunction;
+
+        // Create a function object for module code, similar to CPython's PyEval_EvalCode
+        let qualname = self.ctx.new_str(code.obj_name.as_str());
+        let func = PyFunction::new(
+            code.clone(),
+            scope.globals.clone(),
+            None, // closure
+            None, // defaults
+            None, // kwdefaults
+            qualname,
+            self.ctx.new_tuple(vec![]), // type_params
+            self.ctx.new_dict(),        // annotations
+            self.ctx.none(),            // doc
+            self,
+        )?;
+        let func_obj = func.into_ref(&self.ctx).into();
+
+        let frame = Frame::new(code, scope, self.builtins.dict(), &[], Some(func_obj), self)
+            .into_ref(&self.ctx);
         self.run_frame(frame)
     }
 
