1 files changed, 81 insertions, 52 deletions

diff --git a/Modules/_peg_parser.c b/Modules/_peg_parser.c
index 3b27b2c..b66d5a8 100644
--- a/Modules/_peg_parser.c
+++ b/Modules/_peg_parser.c
@@ -1,104 +1,133 @@
 #include <Python.h>
 #include "pegen_interface.h"
 
-PyObject *
-_Py_parse_file(PyObject *self, PyObject *args, PyObject *kwds)
+static int
+_mode_str_to_int(char *mode_str)
 {
-    static char *keywords[] = {"file", "mode", NULL};
-    char *filename;
-    char *mode_str = "exec";
-
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|s", keywords, &filename, &mode_str)) {
-        return NULL;
-    }
-
     int mode;
     if (strcmp(mode_str, "exec") == 0) {
         mode = Py_file_input;
     }
+    else if (strcmp(mode_str, "eval") == 0) {
+        mode = Py_eval_input;
+    }
     else if (strcmp(mode_str, "single") == 0) {
         mode = Py_single_input;
     }
     else {
-        return PyErr_Format(PyExc_ValueError, "mode must be either 'exec' or 'single'");
+        mode = -1;
     }
+    return mode;
+}
 
-    PyArena *arena = PyArena_New();
-    if (arena == NULL) {
+static mod_ty
+_run_parser(char *str, char *filename, int mode, PyCompilerFlags *flags, PyArena *arena, int oldparser)
+{
+    mod_ty mod;
+    if (!oldparser) {
+        mod = PyPegen_ASTFromString(str, filename, mode, flags, arena);
+    }
+    else {
+        mod = PyParser_ASTFromString(str, filename, mode, flags, arena);
+    }
+    return mod;
+}
+
+PyObject *
+_Py_compile_string(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *keywords[] = {"string", "filename", "mode", "oldparser", NULL};
+    char *the_string;
+    char *filename = "<string>";
+    char *mode_str = "exec";
+    int oldparser = 0;
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|ssp", keywords,
+            &the_string, &filename, &mode_str, &oldparser)) {
         return NULL;
     }
 
+    int mode = _mode_str_to_int(mode_str);
+    if (mode == -1) {
+        return PyErr_Format(PyExc_ValueError, "mode must be either 'exec' or 'eval' or 'single'");
+    }
+
     PyCompilerFlags flags = _PyCompilerFlags_INIT;
-    PyObject *result = NULL;
+    flags.cf_flags = PyCF_IGNORE_COOKIE;
 
-    mod_ty res = PyPegen_ASTFromFilename(filename, mode, &flags, arena);
-    if (res == NULL) {
-        goto error;
+    PyArena *arena = PyArena_New();
+    if (arena == NULL) {
+        return NULL;
+    }
+
+    mod_ty mod = _run_parser(the_string, filename, mode, &flags, arena, oldparser);
+    if (mod == NULL) {
+        PyArena_Free(arena);
+        return NULL;
     }
-    result = PyAST_mod2obj(res);
 
-error:
+    PyObject *filename_ob = PyUnicode_DecodeFSDefault(filename);
+    if (filename_ob == NULL) {
+        PyArena_Free(arena);
+        return NULL;
+    }
+    PyCodeObject *result = PyAST_CompileObject(mod, filename_ob, &flags, -1, arena);
+    Py_XDECREF(filename_ob);
     PyArena_Free(arena);
-    return result;
+    return (PyObject *)result;
 }
 
 PyObject *
 _Py_parse_string(PyObject *self, PyObject *args, PyObject *kwds)
 {
-    static char *keywords[] = {"string", "mode", "oldparser", NULL};
+    static char *keywords[] = {"string", "filename", "mode", "oldparser", NULL};
     char *the_string;
+    char *filename = "<string>";
     char *mode_str = "exec";
     int oldparser = 0;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|sp", keywords,
-            &the_string, &mode_str, &oldparser)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|ssp", keywords,
+            &the_string, &filename, &mode_str, &oldparser)) {
         return NULL;
     }
 
-    int mode;
-    if (strcmp(mode_str, "exec") == 0) {
-        mode = Py_file_input;
-    }
-    else if (strcmp(mode_str, "eval") == 0) {
-        mode = Py_eval_input;
-    }
-    else if (strcmp(mode_str, "single") == 0) {
-        mode = Py_single_input;
-    }
-    else {
+    int mode = _mode_str_to_int(mode_str);
+    if (mode == -1) {
         return PyErr_Format(PyExc_ValueError, "mode must be either 'exec' or 'eval' or 'single'");
     }
 
+    PyCompilerFlags flags = _PyCompilerFlags_INIT;
+    flags.cf_flags = PyCF_IGNORE_COOKIE;
+
     PyArena *arena = PyArena_New();
     if (arena == NULL) {
         return NULL;
     }
 
-    PyObject *result = NULL;
-
-    PyCompilerFlags flags = _PyCompilerFlags_INIT;
-    flags.cf_flags = PyCF_IGNORE_COOKIE;
-
-    mod_ty res;
-    if (oldparser) {
-        res = PyParser_ASTFromString(the_string, "<string>", mode, &flags, arena);
-    }
-    else {
-        res = PyPegen_ASTFromString(the_string, "<string>", mode, &flags, arena);
-    }
-    if (res == NULL) {
-        goto error;
+    mod_ty mod = _run_parser(the_string, filename, mode, &flags, arena, oldparser);
+    if (mod == NULL) {
+        PyArena_Free(arena);
+        return NULL;
     }
-    result = PyAST_mod2obj(res);
 
-error:
+    PyObject *result = PyAST_mod2obj(mod);
     PyArena_Free(arena);
     return result;
 }
 
 static PyMethodDef ParseMethods[] = {
-    {"parse_file", (PyCFunction)(void (*)(void))_Py_parse_file, METH_VARARGS|METH_KEYWORDS, "Parse a file."},
-    {"parse_string", (PyCFunction)(void (*)(void))_Py_parse_string, METH_VARARGS|METH_KEYWORDS,"Parse a string."},
+    {
+        "parse_string",
+        (PyCFunction)(void (*)(void))_Py_parse_string,
+        METH_VARARGS|METH_KEYWORDS,
+        "Parse a string, return an AST."
+    },
+    {
+        "compile_string",
+        (PyCFunction)(void (*)(void))_Py_compile_string,
+        METH_VARARGS|METH_KEYWORDS,
+        "Compile a string, return a code object."
+    },
     {NULL, NULL, 0, NULL} /* Sentinel */
 };