gh-98811: use full source location to simplify __future__ imports error checking. This also fixes an incorrect error offset. (GH-98812)

5 files changed, 71 insertions, 69 deletions

diff --git a/Include/cpython/compile.h b/Include/cpython/compile.h
index 518a376..f5a62a8 100644
--- a/Include/cpython/compile.h
+++ b/Include/cpython/compile.h
@@ -31,11 +31,26 @@ typedef struct {
 #define _PyCompilerFlags_INIT \
     (PyCompilerFlags){.cf_flags = 0, .cf_feature_version = PY_MINOR_VERSION}
 
+/* source location information */
+typedef struct {
+    int lineno;
+    int end_lineno;
+    int col_offset;
+    int end_col_offset;
+} _PyCompilerSrcLocation;
+
+#define SRC_LOCATION_FROM_AST(n) \
+    (_PyCompilerSrcLocation){ \
+               .lineno = (n)->lineno, \
+               .end_lineno = (n)->end_lineno, \
+               .col_offset = (n)->col_offset, \
+               .end_col_offset = (n)->end_col_offset }
+
 /* Future feature support */
 
 typedef struct {
-    int ff_features;      /* flags set by future statements */
-    int ff_lineno;        /* line number of last future statement */
+    int ff_features;                    /* flags set by future statements */
+    _PyCompilerSrcLocation ff_location; /* location of last future statement */
 } PyFutureFeatures;
 
 #define FUTURE_NESTED_SCOPES "nested_scopes"
diff --git a/Lib/test/test_future.py b/Lib/test/test_future.py
index 189cbdc..b8b591a 100644
--- a/Lib/test/test_future.py
+++ b/Lib/test/test_future.py
@@ -60,7 +60,7 @@ class FutureTest(unittest.TestCase):
     def test_badfuture7(self):
         with self.assertRaises(SyntaxError) as cm:
             from test import badsyntax_future7
-        self.check_syntax_error(cm.exception, "badsyntax_future7", 3, 53)
+        self.check_syntax_error(cm.exception, "badsyntax_future7", 3, 54)
 
     def test_badfuture8(self):
         with self.assertRaises(SyntaxError) as cm:
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-10-28-13-59-51.gh-issue-98811.XQypJa.rst b/Misc/NEWS.d/next/Core and Builtins/2022-10-28-13-59-51.gh-issue-98811.XQypJa.rst
new file mode 100644
index 0000000..ce90a51
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-10-28-13-59-51.gh-issue-98811.XQypJa.rst
@@ -0,0 +1,4 @@
+Use complete source locations to simplify detection of ``__future__``
+imports which are not at the beginning of the file. Also corrects the offset
+in the exception raised in one case, which was off by one and impeded
+highlighting.
diff --git a/Python/compile.c b/Python/compile.c
index 7797291..f892478 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -125,18 +125,23 @@
         (c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT) \
         && (c->u->u_ste->ste_type == ModuleBlock))
 
-typedef struct location_ {
-    int lineno;
-    int end_lineno;
-    int col_offset;
-    int end_col_offset;
-} location;
+typedef _PyCompilerSrcLocation location;
 
 #define LOCATION(LNO, END_LNO, COL, END_COL) \
     ((const location){(LNO), (END_LNO), (COL), (END_COL)})
 
 static location NO_LOCATION = {-1, -1, -1, -1};
 
+/* Return true if loc1 starts after loc2 ends. */
+static inline bool
+location_is_after(location loc1, location loc2) {
+    return (loc1.lineno > loc2.end_lineno) ||
+            ((loc1.lineno == loc2.end_lineno) &&
+             (loc1.col_offset > loc2.end_col_offset));
+}
+
+#define LOC(x) SRC_LOCATION_FROM_AST(x)
+
 typedef struct jump_target_label_ {
     int id;
 } jump_target_label;
@@ -1012,11 +1017,6 @@ basicblock_next_instr(basicblock *b)
 // Artificial instructions
 #define UNSET_LOC(c)
 
-#define LOC(x) LOCATION((x)->lineno,          \
-                        (x)->end_lineno,      \
-                        (x)->col_offset,      \
-                        (x)->end_col_offset)
-
 
 /* Return the stack effect of opcode with argument oparg.
 
@@ -3911,59 +3911,61 @@ compiler_import(struct compiler *c, stmt_ty s)
 static int
 compiler_from_import(struct compiler *c, stmt_ty s)
 {
-    location loc = LOC(s);
-    Py_ssize_t i, n = asdl_seq_LEN(s->v.ImportFrom.names);
-    PyObject *names;
+    Py_ssize_t n = asdl_seq_LEN(s->v.ImportFrom.names);
 
-    ADDOP_LOAD_CONST_NEW(c, loc, PyLong_FromLong(s->v.ImportFrom.level));
+    ADDOP_LOAD_CONST_NEW(c, LOC(s), PyLong_FromLong(s->v.ImportFrom.level));
 
-    names = PyTuple_New(n);
-    if (!names)
+    PyObject *names = PyTuple_New(n);
+    if (!names) {
         return 0;
+    }
 
     /* build up the names */
-    for (i = 0; i < n; i++) {
+    for (Py_ssize_t i = 0; i < n; i++) {
         alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, i);
         Py_INCREF(alias->name);
         PyTuple_SET_ITEM(names, i, alias->name);
     }
 
-    if (s->lineno > c->c_future->ff_lineno && s->v.ImportFrom.module &&
-        _PyUnicode_EqualToASCIIString(s->v.ImportFrom.module, "__future__")) {
+    if (location_is_after(LOC(s), c->c_future->ff_location) &&
+        s->v.ImportFrom.module &&
+        _PyUnicode_EqualToASCIIString(s->v.ImportFrom.module, "__future__"))
+    {
         Py_DECREF(names);
-        return compiler_error(c, loc, "from __future__ imports must occur "
+        return compiler_error(c, LOC(s), "from __future__ imports must occur "
                               "at the beginning of the file");
     }
-    ADDOP_LOAD_CONST_NEW(c, loc, names);
+    ADDOP_LOAD_CONST_NEW(c, LOC(s), names);
 
     if (s->v.ImportFrom.module) {
-        ADDOP_NAME(c, loc, IMPORT_NAME, s->v.ImportFrom.module, names);
+        ADDOP_NAME(c, LOC(s), IMPORT_NAME, s->v.ImportFrom.module, names);
     }
     else {
         _Py_DECLARE_STR(empty, "");
-        ADDOP_NAME(c, loc, IMPORT_NAME, &_Py_STR(empty), names);
+        ADDOP_NAME(c, LOC(s), IMPORT_NAME, &_Py_STR(empty), names);
     }
-    for (i = 0; i < n; i++) {
+    for (Py_ssize_t i = 0; i < n; i++) {
         alias_ty alias = (alias_ty)asdl_seq_GET(s->v.ImportFrom.names, i);
         identifier store_name;
 
         if (i == 0 && PyUnicode_READ_CHAR(alias->name, 0) == '*') {
             assert(n == 1);
-            ADDOP(c, loc, IMPORT_STAR);
+            ADDOP(c, LOC(s), IMPORT_STAR);
             return 1;
         }
 
-        ADDOP_NAME(c, loc, IMPORT_FROM, alias->name, names);
+        ADDOP_NAME(c, LOC(s), IMPORT_FROM, alias->name, names);
         store_name = alias->name;
-        if (alias->asname)
+        if (alias->asname) {
             store_name = alias->asname;
+        }
 
-        if (!compiler_nameop(c, loc, store_name, Store)) {
+        if (!compiler_nameop(c, LOC(s), store_name, Store)) {
             return 0;
         }
     }
     /* remove imported module */
-    ADDOP(c, loc, POP_TOP);
+    ADDOP(c, LOC(s), POP_TOP);
     return 1;
 }
 
diff --git a/Python/future.c b/Python/future.c
index d465608..2a45d2e 100644
--- a/Python/future.c
+++ b/Python/future.c
@@ -2,8 +2,6 @@
 #include "pycore_ast.h"           // _PyAST_GetDocString()
 
 #define UNDEFINED_FUTURE_FEATURE "future feature %.100s is not defined"
-#define ERR_LATE_FUTURE \
-"from __future__ imports must occur at the beginning of the file"
 
 static int
 future_check_features(PyFutureFeatures *ff, stmt_ty s, PyObject *filename)
@@ -56,59 +54,42 @@ future_check_features(PyFutureFeatures *ff, stmt_ty s, PyObject *filename)
 static int
 future_parse(PyFutureFeatures *ff, mod_ty mod, PyObject *filename)
 {
-    int i, done = 0, prev_line = 0;
-
-    if (!(mod->kind == Module_kind || mod->kind == Interactive_kind))
+    if (!(mod->kind == Module_kind || mod->kind == Interactive_kind)) {
         return 1;
+    }
 
-    if (asdl_seq_LEN(mod->v.Module.body) == 0)
+    Py_ssize_t n = asdl_seq_LEN(mod->v.Module.body);
+    if (n == 0) {
         return 1;
+    }
 
-    /* A subsequent pass will detect future imports that don't
-       appear at the beginning of the file.  There's one case,
-       however, that is easier to handle here: A series of imports
-       joined by semi-colons, where the first import is a future
-       statement but some subsequent import has the future form
-       but is preceded by a regular import.
-    */
-
-    i = 0;
-    if (_PyAST_GetDocString(mod->v.Module.body) != NULL)
+    Py_ssize_t i = 0;
+    if (_PyAST_GetDocString(mod->v.Module.body) != NULL) {
         i++;
+    }
 
-    for (; i < asdl_seq_LEN(mod->v.Module.body); i++) {
+    for (; i < n; i++) {
         stmt_ty s = (stmt_ty)asdl_seq_GET(mod->v.Module.body, i);
 
-        if (done && s->lineno > prev_line)
-            return 1;
-        prev_line = s->lineno;
-
-        /* The tests below will return from this function unless it is
-           still possible to find a future statement.  The only things
-           that can precede a future statement are another future
-           statement and a doc string.
-        */
+        /* The only things that can precede a future statement
+         *  are another future statement and a doc string.
+         */
 
         if (s->kind == ImportFrom_kind) {
             identifier modname = s->v.ImportFrom.module;
             if (modname &&
                 _PyUnicode_EqualToASCIIString(modname, "__future__")) {
-                if (done) {
-                    PyErr_SetString(PyExc_SyntaxError,
-                                    ERR_LATE_FUTURE);
-                    PyErr_SyntaxLocationObject(filename, s->lineno, s->col_offset);
+                if (!future_check_features(ff, s, filename)) {
                     return 0;
                 }
-                if (!future_check_features(ff, s, filename))
-                    return 0;
-                ff->ff_lineno = s->lineno;
+                ff->ff_location = SRC_LOCATION_FROM_AST(s);
             }
             else {
-                done = 1;
+                return 1;
             }
         }
         else {
-            done = 1;
+            return 1;
         }
     }
     return 1;
@@ -126,7 +107,7 @@ _PyFuture_FromAST(mod_ty mod, PyObject *filename)
         return NULL;
     }
     ff->ff_features = 0;
-    ff->ff_lineno = -1;
+    ff->ff_location = (_PyCompilerSrcLocation){-1, -1, -1, -1};
 
     if (!future_parse(ff, mod, filename)) {
         PyObject_Free(ff);