Support compile(source, filename, 'func_type', PyCF_ONLY_AST)

diff --git a/Include/compile.h b/Include/compile.h
index 36477c01e2ba6324c98cecc8a6c2159763b78e6f..d0bbed8f558b15288a071e17d67d0bfdf05e72e6 100644 (file)
--- a/Include/compile.h
+++ b/Include/compile.h
@@ -86,10 +86,10 @@ PyAPI_FUNC(int) _PyAST_Optimize(struct _mod *, PyArena *arena, int optimize);
 
 #endif /* !Py_LIMITED_API */
 
-/* These definitions must match corresponding definitions in graminit.h.
-   There's code in compile.c that checks that they are the same. */
+/* These definitions must match corresponding definitions in graminit.h. */
 #define Py_single_input 256
 #define Py_file_input 257
 #define Py_eval_input 258
+#define Py_func_type_input 345
 
 #endif /* !Py_COMPILE_H */

diff --git a/Lib/test/test_type_comments.py b/Lib/test/test_type_comments.py
index 5ed60c35a2f5f5b6634363e1de2cc6911a7a4a3b..e2cc88d701e96470ee8bd8e95ed5cf0bcf2cceb0 100644 (file)
--- a/Lib/test/test_type_comments.py
+++ b/Lib/test/test_type_comments.py
@@ -239,6 +239,25 @@ class TypeCommentTests(unittest.TestCase):
         check_both_ways("try:  # type: int\n  pass\nfinally:\n  pass\n")
         check_both_ways("try:\n  pass\nfinally:  # type: int\n  pass\n")
 
+    def test_func_type_input(self):
+        return
+
+        def parse_func_type_input(source):
+            from ast import PyCF_ONLY_AST
+            return compile(source, "<unknown>", "func_type", PyCF_ONLY_AST)
+
+        # Some checks below will crash if the return structure is wrong
+        tree = parse_func_type_input("() -> int")
+        self.assertEqual(tree.argtypes, [])
+        self.assertEqual(tree.returns.id, "int")
+
+        tree = parse_func_type_input("(int) -> List[str]")
+        self.assertEqual(len(tree.argtypes), 1)
+        arg = tree.argtypes[0]
+        self.assertEqual(arg.value.int, "int")
+        self.assertEqual(tree.returns.value.id, "List")
+        self.assertEqual(tree.returns.slice.vale.id, "str")
+
 
 if __name__ == '__main__':
     unittest.main()

diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 5467d192ee69c07443bd39150df550dcca6842f1..1a56e90bca09aa09d61a589cfbc1b56c36be4aaf 100644 (file)
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -8917,18 +8917,19 @@ PyObject* PyAST_mod2obj(mod_ty t)
 }
 
 /* mode is 0 for "exec", 1 for "eval" and 2 for "single" input */
+/* and 3 for "func_type" */
 mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
 {
     mod_ty res;
     PyObject *req_type[3];
-    char *req_name[] = {"Module", "Expression", "Interactive"};
+    char *req_name[] = {"Module", "Expression", "Interactive", "FunctionType"};
     int isinstance;
 
     req_type[0] = (PyObject*)Module_type;
     req_type[1] = (PyObject*)Expression_type;
     req_type[2] = (PyObject*)Interactive_type;
 
-    assert(0 <= mode && mode <= 2);
+    assert(0 <= mode && mode <= 3);
 
     if (!init_types())
         return NULL;

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 930adb496e14cd7aaccffdba175992e8b75c4867..f9b901f7e59fb99d647550ba24e71e0d7eec4bc2 100644 (file)
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -765,7 +765,7 @@ builtin_compile_impl(PyObject *module, PyObject *source, PyObject *filename,
     int compile_mode = -1;
     int is_ast;
     PyCompilerFlags cf;
-    int start[] = {Py_file_input, Py_eval_input, Py_single_input};
+    int start[] = {Py_file_input, Py_eval_input, Py_single_input, Py_func_type_input};
     PyObject *result;
 
     cf.cf_flags = flags | PyCF_SOURCE_IS_UTF8;
@@ -795,9 +795,21 @@ builtin_compile_impl(PyObject *module, PyObject *source, PyObject *filename,
         compile_mode = 1;
     else if (strcmp(mode, "single") == 0)
         compile_mode = 2;
+    else if (strcmp(mode, "func_type") == 0) {
+        if (!(flags & PyCF_ONLY_AST)) {
+            PyErr_SetString(PyExc_ValueError,
+                            "compile() mode 'func_type' requires flag PyCF_ONLY_AST");
+            goto error;
+        }
+        compile_mode = 3;
+    }
     else {
-        PyErr_SetString(PyExc_ValueError,
-                        "compile() mode must be 'exec', 'eval' or 'single'");
+        const char *msg;
+        if (flags & PyCF_ONLY_AST)
+            msg = "compile() mode must be 'exec', 'eval', 'single' or 'func_type'";
+        else
+            msg = "compile() mode must be 'exec', 'eval' or 'single'";
+        PyErr_SetString(PyExc_ValueError, msg);
         goto error;
     }