bpo-41631: Check that importing _ast returns the right module

Lib/test/test_ast.py

@@ -1945,6 +1945,22 @@ def visit_Ellipsis(self, node):
         ])
 
 
+class ModuleReplacementTests(unittest.TestCase):
+    # bpo-41631
+    # Check that compile(..., ast.PyCF_ONLY_AST) fails if _ast is tampered with
+    def test_replaced_ast_module(self):
+        old_ast = sys.modules['_ast']
+        try:
+            for replacement in None, 'a string', ast:
+                with self.subTest(replacement=replacement):
+                    sys.modules['_ast'] = replacement
+                    with self.assertRaises(SystemError):
+                        compile('1+1', '<string>', 'eval',
+                                flags=ast.PyCF_ONLY_AST)
+        finally:
+            sys.modules['_ast'] = old_ast
+
+
 def main():
     if __name__ != '__main__':
         return

Misc/NEWS.d/next/Core and Builtins/2020-08-26-22-25-25.bpo-41631.RTN8yW.rst

@@ -0,0 +1,2 @@
+If the _ast module is unavailable or replaced with an unexpected object,
+compiling to AST now raises an exception.

Parser/asdl_c.py

@@ -1382,6 +1382,8 @@ def generate_module_def(f, mod):
     return (astmodulestate*)state;
 }
 
+static struct PyModuleDef _astmodule;
+
 static astmodulestate*
 get_global_ast_state(void)
 {
@@ -1400,6 +1402,29 @@ def generate_module_def(f, mod):
             return NULL;
         }
     }
+    if (!PyModule_Check(module)) {
+        PyErr_SetString(
+            PyExc_SystemError,
+            "Non-module object imported as _ast");
+        return NULL;
+    }
+
+// testing "pegen" requires generating a C extension module, which contains
+// a copy of the symbols defined in Python-ast.c and needs to be treated
+// as the _ast module.
+#ifndef AST_SKIP_MODULE_CHECK
+
+    struct PyModuleDef *def = PyModule_GetDef(module);
+    if (def != &_astmodule) {
+        if (!PyErr_Occurred()) {
+            PyErr_SetString(
+                PyExc_SystemError,
+                "Unexpected module imported as _ast");
+        }
+        return NULL;
+    }
+#endif
+
     astmodulestate *state = get_ast_state(module);
     Py_DECREF(module);
     return state;

Tools/peg_generator/pegen/ast_dump.py

@@ -3,7 +3,8 @@
 because testing pegen requires generating a C extension module, which contains
 a copy of the symbols defined in Python-ast.c. Thus, the isinstance check would
 always fail. We rely on string comparison of the base classes instead.
-TODO: Remove the above-described hack.
+TODO: Remove the above-described hack, along with AST_SKIP_MODULE_CHECK
+elsewhere.
 """
 
 

Tools/peg_generator/pegen/build.py

@@ -59,6 +59,12 @@ def compile_c_extension(
     extra_link_args = get_extra_flags("LDFLAGS", "PY_LDFLAGS_NODIST")
     if keep_asserts:
         extra_compile_args.append("-UNDEBUG")
+
+    # testing pegen requires generating C extension modules, which contains
+    # a copy of the symbols defined in Python-ast.c and needs to be treated
+    # as the _ast module. See ast_dump.py.
+    extra_compile_args.append("-DAST_SKIP_MODULE_CHECK")
+
     extension = [
         Extension(
             extension_name,