diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 1f290d442bb..fc115c60977 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -2996,8 +2996,7 @@ impl<'warnings> Compiler<'warnings> { self.emit_load_const(ConstantData::Str { value: value.into(), }); - let doc = self.name("__doc__"); - emit!(self, Instruction::StoreName { namei: doc }); + self.store_name(&"__doc__".into())?; self.set_no_location(); self.set_source_range(saved_range); } @@ -20054,6 +20053,54 @@ class C: })); } + fn find_module_doc_store(code: &CodeObject) -> Option<&Instruction> { + code.instructions.iter().find_map(|unit| match unit.op { + Instruction::StoreName { namei } | Instruction::StoreGlobal { namei } => { + let idx = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + (code.names[usize::try_from(idx).unwrap()].as_str() == "__doc__") + .then_some(&unit.op) + } + _ => None, + }) + } + + #[test] + fn module_docstring_store_uses_global_when_nested_scope_declares_global() { + let code = compile_exec( + "\ +\"module doc\" +def f(): + global __doc__ +", + ); + + let op = find_module_doc_store(&code).expect("missing module docstring store"); + assert!( + matches!(op, Instruction::StoreGlobal { .. }), + "module docstring store should resolve through name resolution to StoreGlobal \ + when a nested scope declares `global __doc__`, got instructions={:?}", + code.instructions + ); + } + + #[test] + fn module_docstring_store_stays_name_without_nested_global() { + let code = compile_exec( + "\ +\"module doc\" +x = 1 +", + ); + + let op = find_module_doc_store(&code).expect("missing module docstring store"); + assert!( + matches!(op, Instruction::StoreName { .. }), + "module docstring store should remain StoreName without a nested \ + `global __doc__` declaration, got instructions={:?}", + code.instructions + ); + } + #[test] fn conditional_return_epilogue_is_duplicated() { let code = compile_exec( diff --git a/extra_tests/snippets/syntax_doc.py b/extra_tests/snippets/syntax_doc.py index bdfc1fe7788..37aebb43dbd 100644 --- a/extra_tests/snippets/syntax_doc.py +++ b/extra_tests/snippets/syntax_doc.py @@ -13,3 +13,77 @@ def f2(): """ assert f2.__doc__ == '\nx\n y\n' + + +# Regression test: a module docstring must be stored through module name +# resolution, so a nested `global __doc__` declaration makes the module body +# emit STORE_GLOBAL (writing module globals) instead of STORE_NAME (writing +# the provided locals mapping). Mirrors CPython codegen_nameop behavior. + +import dis + + +def opnames(code): + return [instr.opname for instr in dis.get_instructions(code)] + + +def run_split(src, optimize=0): + code = compile(src, "", "exec", optimize=optimize) + g, loc = {}, {} + exec(code, g, loc) + return code, g, loc + + +# A nested function declaring `global __doc__` makes the module docstring +# store target module globals. +src_global = '"module doc"\ndef f():\n global __doc__\n' +code, g, loc = run_split(src_global) +ops = opnames(code) +assert "STORE_GLOBAL" in ops, ops +assert "STORE_NAME" not in ops or "__doc__" not in [ + instr.argval for instr in dis.get_instructions(code) if instr.opname == "STORE_NAME" +], ops +assert g.get("__doc__") == "module doc", g +assert "__doc__" not in loc, loc + +# The declaration may be nested more than one scope deep. +src_deep = '"module doc"\ndef f():\n def g():\n global __doc__\n' +code, g, loc = run_split(src_deep) +assert "STORE_GLOBAL" in opnames(code), opnames(code) +assert g.get("__doc__") == "module doc", g +assert "__doc__" not in loc, loc + +# Control: without a nested global declaration the docstring stays a +# STORE_NAME and lands in the provided locals mapping. +src_plain = '"module doc"\nx = 1\n' +code, g, loc = run_split(src_plain) +doc_stores = [ + instr.opname + for instr in dis.get_instructions(code) + if instr.argval == "__doc__" and instr.opname.startswith("STORE_") +] +assert doc_stores == ["STORE_NAME"], doc_stores +assert loc.get("__doc__") == "module doc", loc +assert "__doc__" not in g, g + +# optimize=1 keeps the docstring and the same global resolution. +code, g, loc = run_split(src_global, optimize=1) +assert "STORE_GLOBAL" in opnames(code), opnames(code) +assert g.get("__doc__") == "module doc", g + +# optimize=2 strips the docstring entirely: no __doc__ store in either dict. +code, g, loc = run_split(src_global, optimize=2) +assert not any( + instr.argval == "__doc__" and instr.opname.startswith("STORE_") + for instr in dis.get_instructions(code) +), opnames(code) +assert "__doc__" not in g and "__doc__" not in loc, (g, loc) + +# A class docstring still belongs to the class namespace even when a nested +# scope declares `global __doc__`. +src_class = '"module doc"\ndef f():\n global __doc__\nclass C:\n "class doc"\n' +g, loc = {}, {} +exec(compile(src_class, "", "exec"), g, loc) +assert g.get("__doc__") == "module doc", g +assert loc["C"].__doc__ == "class doc", loc["C"].__doc__ +assert "__doc__" not in loc, loc