From 895a0129f1937637a5a0128cc5352d9c096ab3d9 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@python.org>
Date: Mon, 28 Jan 2019 08:57:56 -0800
Subject: [PATCH] Support compile(source, filename, 'func_type', PyCF_ONLY_AST)

---
 Include/compile.h              |  4 ++--
 Lib/test/test_type_comments.py | 19 +++++++++++++++++++
 Python/Python-ast.c            |  5 +++--
 Python/bltinmodule.c           | 18 +++++++++++++++---
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/Include/compile.h b/Include/compile.h
index 36477c01e2ba..d0bbed8f558b 100644
--- 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 5ed60c35a2f5..e2cc88d701e9 100644
--- 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 5467d192ee69..1a56e90bca09 100644
--- 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 930adb496e14..f9b901f7e59f 100644
--- 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;
     }
 
-- 
2.47.3