Apply suggestions from code review

pablogsal · lysnikolaou · nineteendo · pablogsal · commit cd167c2153ab · 2024-10-29T23:27:50.000Z
Co-authored-by: Lysandros Nikolaou &lt;lisandrosnik@gmail.com&gt;
Co-authored-by: Nice Zombies &lt;nineteendo19d0@gmail.com&gt;
Signed-off-by: Pablo Galindo &lt;pablogsal@gmail.com&gt;
diff --git a/Grammar/python.gram b/Grammar/python.gram
@@ -207,9 +207,7 @@ import_name[stmt_ty]: 'import' a=dotted_as_names { _PyAST_Import(a, EXTRA) }
 # note below: the ('.' | '...') is necessary because '...' is tokenized as ELLIPSIS
 import_from[stmt_ty]:
     | 'from' a=('.' | '...')* b=dotted_name 'import' c=import_from_targets {
-        _PyPegen_check_future_import(p,
-            _PyAST_ImportFrom(b->v.Name.id, c, _PyPegen_seq_count_dots(a), EXTRA)
-        ) }
+        _PyPegen_checked_future_import(p, b->v.Name.id, c, _PyPegen_seq_count_dots(a), EXTRA) }
     | 'from' a=('.' | '...')+ 'import' b=import_from_targets {
         _PyAST_ImportFrom(NULL, b, _PyPegen_seq_count_dots(a), EXTRA) }
 import_from_targets[asdl_alias_seq*]:
diff --git a/Lib/test/test_grammar.py b/Lib/test/test_grammar.py
@@ -1972,6 +1972,12 @@ async def foo():
         with self.assertRaises(Done):
             foo().send(None)
 
+    def test_inline_future_imports(self):
+        # Check that future imports that change the behavior of the parser are
+        # correctly handled by the parser.
+        compile("from __future__ import barry_as_FLUFL; 1<>2", "<string>", "exec")
+        compile("from __future__ import annotations, barry_as_FLUFL; 1<>2", "<string>", "exec")
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/Parser/action_helpers.c b/Parser/action_helpers.c
@@ -1693,22 +1693,16 @@ _PyPegen_concatenate_strings(Parser *p, asdl_expr_seq *strings,
 }
 
 stmt_ty
-_PyPegen_check_future_import(Parser *p, stmt_ty importfrom) {
-    if (importfrom->kind != ImportFrom_kind) {
-        goto exit;
-    }
-    asdl_alias_seq *names = importfrom->v.ImportFrom.names;
-    if (asdl_seq_LEN(names) != 1) {
-        goto exit;
-    }
-    identifier mod = importfrom->v.ImportFrom.module;
-    if (PyUnicode_CompareWithASCIIString(mod, "__future__") != 0) {
-        goto exit;
-    }
-    alias_ty alias = asdl_seq_GET(names, 0);
-    if (PyUnicode_CompareWithASCIIString(alias->name, "barry_as_FLUFL") == 0) {
-        p->flags |= PyPARSE_BARRY_AS_BDFL;
+_PyPegen_checked_future_import(Parser *p, identifier module, asdl_alias_seq * names, int level,
+                  int lineno, int col_offset, int end_lineno, int end_col_offset,
+                  PyArena *arena) {
+    if (PyUnicode_CompareWithASCIIString(module, "__future__") == 0) {
+        for (Py_ssize_t i = 0; i < asdl_seq_LEN(names); i++) {
+            alias_ty alias = asdl_seq_GET(names, i);
+            if (PyUnicode_CompareWithASCIIString(alias->name, "barry_as_FLUFL") == 0) {
+                p->flags |= PyPARSE_BARRY_AS_BDFL;
+            }
+        }
     }
-exit:
-    return importfrom;
+    return _PyAST_ImportFrom(module, names, level, lineno, col_offset, end_lineno, end_col_offset, arena);
 }
diff --git a/Parser/parser.c b/Parser/parser.c
diff --git a/Parser/pegen.h b/Parser/pegen.h
@@ -346,6 +346,8 @@ mod_ty _PyPegen_make_module(Parser *, asdl_stmt_seq *);
 void *_PyPegen_arguments_parsing_error(Parser *, expr_ty);
 expr_ty _PyPegen_get_last_comprehension_item(comprehension_ty comprehension);
 void *_PyPegen_nonparen_genexp_in_call(Parser *p, expr_ty args, asdl_comprehension_seq *comprehensions);
+stmt_ty _PyPegen_checked_future_import(Parser *p, identifier module, asdl_alias_seq *,
+                                       int , int, int , int , int , PyArena *);
 
 // Parser API
 
@@ -361,8 +363,6 @@ asdl_stmt_seq *_PyPegen_interactive_exit(Parser *);
 // TODO: move to the correct place in this file
 expr_ty _PyPegen_joined_str(Parser *p, Token* a, asdl_expr_seq* expr, Token*b);
 
-stmt_ty _PyPegen_check_future_import(Parser *p, stmt_ty importfrom);
-
 // Generated function in parse.c - function definition in python.gram
 void *_PyPegen_parse(Parser *);
 
