1 files changed, 25 insertions, 0 deletions
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index 721c527745c..b2e04cfa317 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -21,6 +21,7 @@
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException, _Py_Offer_Suggestions
 #include "pycore_pylifecycle.h"   // _Py_UnhandledKeyboardInterrupt
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
+#include "pycore_symtable.h"      // _PyFuture_FromAST()
 #include "pycore_sysmodule.h"     // _PySys_Audit()
 #include "pycore_traceback.h"     // _PyTraceBack_Print_Indented()
 
@@ -1790,6 +1791,24 @@ error:
     return NULL;
 }
 
+static int
+ast_optimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
+             int optimize, PyArena *arena)
+{
+    PyFutureFeatures future;
+    if (!_PyFuture_FromAST(mod, filename, &future)) {
+        return -1;
+    }
+    int flags = future.ff_features | cf->cf_flags;
+    if (optimize == -1) {
+        optimize = _Py_GetConfig()->optimization_level;
+    }
+    if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
+        return -1;
+    }
+    return 0;
+}
+
 PyObject *
 Py_CompileStringObject(const char *str, PyObject *filename, int start,
                        PyCompilerFlags *flags, int optimize)
@@ -1806,6 +1825,12 @@ Py_CompileStringObject(const char *str, PyObject *filename, int start,
         return NULL;
     }
     if (flags && (flags->cf_flags & PyCF_ONLY_AST)) {
+        if ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_OPTIMIZED_AST) {
+            if (ast_optimize(mod, filename, flags, optimize, arena) < 0) {
+                _PyArena_Free(arena);
+                return NULL;
+            }
+        }
         PyObject *result = PyAST_mod2obj(mod);
         _PyArena_Free(arena);
         return result;