bpo-43950: support positions for dis.Instructions created through dis.Bytecode (GH-28142)

isidentical · isidentical · commit 6ed3a9d25fe2 · 2021-09-09T00:46:57.000+03:00
diff --git a/Parser/pegen.c b/Parser/pegen.c
@@ -410,7 +410,54 @@ get_error_line(Parser *p, Py_ssize_t lineno)
 }
 
 Py_ssize_t
-_PyPegen_byte_offset_to_character_offset(PyObject *line, Py_ssize_t col_offset)
+byte_offset_difference_for_unicode(Py_ssize_t size, PyObject *text)
+{
+    if (PyUnicode_IS_ASCII(text)) {
+        return 0;
+    }
+
+    PyObject *unicode_data = PyImport_ImportModuleNoBlock("unicodedata");
+    if (!unicode_data) {
+        return -1;
+    }
+
+    PyObject *east_asian_width = PyObject_GetAttrString(unicode_data, "east_asian_width");
+    Py_DECREF(unicode_data);
+    if (!east_asian_width) {
+        return -1;
+    }
+
+    PyObject *current, *width = NULL;
+    Py_ssize_t i, difference = 0;
+    for (i = 0; i < size; i++) {
+        current = PySequence_GetItem(text, i);
+        if (!current) {
+            goto error;
+        }
+
+        PyObject *width = PyObject_CallOneArg(east_asian_width, current);
+        Py_DECREF(current);
+        if (!width) {
+            goto error;
+        }
+
+        if (_PyUnicode_EqualToASCIIString(width, "W") ||
+            _PyUnicode_EqualToASCIIString(width, "F")) {
+            ++difference;
+        }
+        Py_DECREF(width);
+    }
+    return difference;
+
+error:
+    Py_DECREF(east_asian_width);
+    Py_XDECREF(current);
+    Py_XDECREF(width);
+    return -1;
+}
+
+Py_ssize_t
+_PyPegen_byte_offset_to_character_offset(PyObject *line, Py_ssize_t col_offset, Py_ssize_t *difference)
 {
     const char *str = PyUnicode_AsUTF8(line);
     if (!str) {
@@ -425,7 +472,10 @@ _PyPegen_byte_offset_to_character_offset(PyObject *line, Py_ssize_t col_offset)
     if (!text) {
         return -1;
     }
+
     Py_ssize_t size = PyUnicode_GET_LENGTH(text);
+    *difference = byte_offset_difference_for_unicode(size, text);
+
     Py_DECREF(text);
     return size;
 }
@@ -508,14 +558,15 @@ _PyPegen_raise_error_known_location(Parser *p, PyObject *errtype,
 
     Py_ssize_t col_number = col_offset;
     Py_ssize_t end_col_number = end_col_offset;
+    Py_ssize_t start_difference, end_difference;
 
     if (p->tok->encoding != NULL) {
-        col_number = _PyPegen_byte_offset_to_character_offset(error_line, col_offset);
+        col_number = _PyPegen_byte_offset_to_character_offset(error_line, col_offset, &start_difference);
         if (col_number < 0) {
             goto error;
         }
         if (end_col_number > 0) {
-            Py_ssize_t end_col_offset = _PyPegen_byte_offset_to_character_offset(error_line, end_col_number);
+            Py_ssize_t end_col_offset = _PyPegen_byte_offset_to_character_offset(error_line, end_col_number, &end_difference);
             if (end_col_offset < 0) {
                 goto error;
             } else {
diff --git a/Parser/pegen.h b/Parser/pegen.h
@@ -140,7 +140,7 @@ expr_ty _PyPegen_name_token(Parser *p);
 expr_ty _PyPegen_number_token(Parser *p);
 void *_PyPegen_string_token(Parser *p);
 const char *_PyPegen_get_expr_name(expr_ty);
-Py_ssize_t _PyPegen_byte_offset_to_character_offset(PyObject *line, Py_ssize_t col_offset);
+Py_ssize_t _PyPegen_byte_offset_to_character_offset(PyObject *line, Py_ssize_t col_offset, Py_ssize_t *difference);
 void *_PyPegen_raise_error(Parser *p, PyObject *errtype, const char *errmsg, ...);
 void *_PyPegen_raise_error_known_location(Parser *p, PyObject *errtype,
                                           Py_ssize_t lineno, Py_ssize_t col_offset,
diff --git a/Python/traceback.c b/Python/traceback.c
@@ -594,8 +594,9 @@ extract_anchors_from_stmt(const char *segment_str, stmt_ty statement, Py_ssize_t
 static int
 extract_anchors_from_line(PyObject *filename, PyObject *line,
                           Py_ssize_t start_offset, Py_ssize_t end_offset,
-                          Py_ssize_t *left_anchor, Py_ssize_t *right_anchor,
-                          char** primary_error_char, char** secondary_error_char)
+                          Py_ssize_t start_difference, Py_ssize_t *left_anchor,
+                          Py_ssize_t *right_anchor, char** primary_error_char,
+                          char** secondary_error_char)
 {
     int res = -1;
     PyArena *arena = NULL;
@@ -638,11 +639,21 @@ extract_anchors_from_line(PyObject *filename, PyObject *line,
         res = 0;
     }
 
-done:
     if (res > 0) {
-        *left_anchor += start_offset;
-        *right_anchor += start_offset;
+        Py_ssize_t left_difference, right_difference;
+
+        *left_anchor = _PyPegen_byte_offset_to_character_offset(segment, *left_anchor, &left_difference);
+        *right_anchor = _PyPegen_byte_offset_to_character_offset(segment, *right_anchor, &right_difference);
+
+        if (*left_anchor < 0 || *right_anchor < 0) {
+            res = -1;
+            goto done;
+        }
+
+        *left_anchor += start_offset + start_difference + left_difference;
+        *right_anchor += start_offset + start_difference + right_difference;
     }
+done:
     Py_XDECREF(segment);
     if (arena) {
         _PyArena_Free(arena);
@@ -665,7 +676,7 @@ ignore_source_errors(void) {
 
 static inline int
 print_error_location_carets(PyObject *f, int offset, Py_ssize_t start_offset, Py_ssize_t end_offset,
-                            Py_ssize_t right_start_offset, Py_ssize_t left_end_offset,
+                            Py_ssize_t left_end_offset, Py_ssize_t right_start_offset,
                             const char *primary, const char *secondary) {
     int err = 0;
     int special_chars = (left_end_offset != -1 || right_start_offset != -1);
@@ -744,13 +755,15 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
 
     // Convert the utf-8 byte offset to the actual character offset so we print the right number of carets.
     assert(source_line);
-    Py_ssize_t start_offset = _PyPegen_byte_offset_to_character_offset(source_line, start_col_byte_offset);
+
+    Py_ssize_t start_difference, end_difference;
+    Py_ssize_t start_offset = _PyPegen_byte_offset_to_character_offset(source_line, start_col_byte_offset, &start_difference);
     if (start_offset < 0) {
         err = ignore_source_errors() < 0;
         goto done;
     }
 
-    Py_ssize_t end_offset = _PyPegen_byte_offset_to_character_offset(source_line, end_col_byte_offset);
+    Py_ssize_t end_offset = _PyPegen_byte_offset_to_character_offset(source_line, end_col_byte_offset, &end_difference);
     if (end_offset < 0) {
         err = ignore_source_errors() < 0;
         goto done;
@@ -764,6 +777,7 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
 
     if (start_line == end_line) {
         int res = extract_anchors_from_line(filename, source_line, start_offset, end_offset,
+                                            start_difference,
                                             &left_end_offset, &right_start_offset,
                                             &primary_error_char, &secondary_error_char);
         if (res < 0 && ignore_source_errors() < 0) {
@@ -788,8 +802,11 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
         end_offset = i + 1;
     }
 
+    start_offset += start_difference;
+    end_offset += end_difference;
+
     err = print_error_location_carets(f, truncation, start_offset, end_offset,
-                                      right_start_offset, left_end_offset,
+                                      left_end_offset, right_start_offset,
                                       primary_error_char, secondary_error_char);
 
 done:
