Update on "Fix Python 3.8 expecttest machinery again, this time for good."

ezyang · ezyang · commit 86da36a3624e · 2021-06-16T12:01:07.000-07:00
In #59709 I attempted to fix the expecttest machinery to work in Python 3.8. However, I noticed that it would fail to do substitutions in this case: ``` self.assertExpectedInline( foo(), """bar""" ) ``` This is because the triple quoted string is not on the same line as the backtrace line number (at the very beginning), and for safety reasons the preexisting regex refused to search beyond the first line. This wasn't a big deal prior to Python 3.8 because the flipped version of the regex simply required the triple quoted string to be flush with the end of the statement (which it typically was!) But it is a big deal now that we only have the start of the statement. I couldn't think of a way to fix this in the current model, so I decided to call in the big guns. Instead of trying to do the regex with only the start xor end line number, I now require you provide BOTH line numbers, and we will only regex within this range. The way we compute these line numbers is by parsing the Python test file with ast, and then searching through statements until we find one that is consistent with the line number reported by the backtrace. If we don't find anything, we conservatively assume that the string lies exactly in the backtrace (and you'll probably fail the substitution in that case.) The resulting code is quite a lot simpler (no more reversed regex) and hopefully more robust, although I suppose we are going to have to do some field testing. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Differential Revision: [D29146943](https://our.internmc.facebook.com/intern/diff/D29146943) [ghstack-poisoned]
diff --git a/torch/testing/_internal/expecttest.py b/torch/testing/_internal/expecttest.py
@@ -183,7 +183,12 @@ def replace_string_literal(src : str, start_lineno : int, end_lineno : int,
     if delta[0] > 0:
         delta[0] += 1  # handle the extra \\\n
 
-    def compute_raw_new_body_and_adjust_delta(m):
+    assert start_lineno <= end_lineno
+    start = nth_line(src, start_lineno)
+    end = nth_eol(src, end_lineno)
+    assert start <= end
+
+    def replace(m):
         s = new_string
         raw = m.group('raw') == 'r'
         if not raw or not ok_for_raw_triple_quoted_string(s, quote=m.group('quote')[0]):
@@ -196,13 +201,6 @@ def compute_raw_new_body_and_adjust_delta(m):
 
         new_body = "\\\n" + s if "\n" in s and not raw else s
         delta[0] -= m.group('body').count("\n")
-        return raw, new_body
-
-    start = nth_line(src, start_lineno)
-    end = nth_eol(src, end_lineno)
-
-    def replace(m):
-        raw, new_body = compute_raw_new_body_and_adjust_delta(m)
         return ''.join(['r' if raw else '',
                         m.group('quote'),
                         new_body,
@@ -246,14 +244,16 @@ def assertExpectedInline(self, actual, expect, skip=0):
                     start_lineno = lineno
                     end_lineno = lineno
                     # Try to give a more accurate bounds based on AST
+                    # NB: this walk is in no specified order (in practice it's
+                    # breadth first)
                     for n in ast.walk(old_ast):
                         if isinstance(n, ast.Expr):
                             if hasattr(n, 'end_lineno'):
                                 assert LINENO_AT_START
-                                end_lineno = n.end_lineno  # type: ignore
-                                break
+                                if n.lineno == start_lineno:
+                                    end_lineno = n.end_lineno  # type: ignore[attr-defined]
                             else:
-                                if n.lineno < end_lineno:
+                                if n.lineno == end_lineno:
                                     start_lineno = n.lineno
 
                     new, delta = replace_string_literal(old, start_lineno, end_lineno, actual)
