import re import unittest import traceback import os import string ACCEPT = os.getenv('EXPECTTEST_ACCEPT') def nth_line(src, lineno): """ Compute the starting index of the n-th line (where n is 1-indexed) >>> nth_line("aaa\\nbb\\nc", 2) 4 """ assert lineno >= 1 pos = 0 for _ in range(lineno - 1): pos = src.find('\n', pos) + 1 return pos def nth_eol(src, lineno): """ Compute the ending index of the n-th line (before the newline, where n is 1-indexed) >>> nth_eol("aaa\\nbb\\nc", 2) 6 """ assert lineno >= 1 pos = -1 for _ in range(lineno): pos = src.find('\n', pos + 1) if pos == -1: return len(src) return pos def normalize_nl(t): return t.replace('\r\n', '\n').replace('\r', '\n') def escape_trailing_quote(s, quote): if s and s[-1] == quote: return s[:-1] + '\\' + quote else: return s class EditHistory(object): def __init__(self): self.state = {} def adjust_lineno(self, fn, lineno): if fn not in self.state: return lineno for edit_loc, edit_diff in self.state[fn]: if lineno > edit_loc: lineno += edit_diff return lineno def seen_file(self, fn): return fn in self.state def record_edit(self, fn, lineno, delta): self.state.setdefault(fn, []).append((lineno, delta)) EDIT_HISTORY = EditHistory() def ok_for_raw_triple_quoted_string(s, quote): """ Is this string representable inside a raw triple-quoted string? Due to the fact that backslashes are always treated literally, some strings are not representable. >>> ok_for_raw_triple_quoted_string("blah", quote="'") True >>> ok_for_raw_triple_quoted_string("'", quote="'") False >>> ok_for_raw_triple_quoted_string("a ''' b", quote="'") False """ return quote * 3 not in s and (not s or s[-1] not in [quote, '\\']) # This operates on the REVERSED string (that's why suffix is first) RE_EXPECT = re.compile(r"^(?P[^\n]*?)" r"(?P'''|" r'""")' r"(?P.*?)" r"(?P=quote)" r"(?Pr?)", re.DOTALL) def replace_string_literal(src, lineno, new_string): r""" Replace a triple quoted string literal with new contents. Only handles printable ASCII correctly at the moment. This will preserve the quote style of the original string, and makes a best effort to preserve raw-ness (unless it is impossible to do so.) Returns a tuple of the replaced string, as well as a delta of number of lines added/removed. >>> replace_string_literal("'''arf'''", 1, "barf") ("'''barf'''", 0) >>> r = replace_string_literal(" moo = '''arf'''", 1, "'a'\n\\b\n") >>> print(r[0]) moo = '''\ 'a' \\b ''' >>> r[1] 3 >>> replace_string_literal(" moo = '''\\\narf'''", 2, "'a'\n\\b\n")[1] 2 >>> print(replace_string_literal(" f('''\"\"\"''')", 1, "a ''' b")[0]) f('''a \'\'\' b''') """ # Haven't implemented correct escaping for non-printable characters assert all(c in string.printable for c in new_string) i = nth_eol(src, lineno) new_string = normalize_nl(new_string) delta = [new_string.count("\n")] if delta[0] > 0: delta[0] += 1 # handle the extra \\\n 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]): raw = False s = s.replace('\\', '\\\\') if m.group('quote') == "'''": s = escape_trailing_quote(s, "'").replace("'''", r"\'\'\'") else: s = escape_trailing_quote(s, '"').replace('"""', r'\"\"\"') new_body = "\\\n" + s if "\n" in s and not raw else s delta[0] -= m.group('body').count("\n") return ''.join([m.group('suffix'), m.group('quote'), new_body[::-1], m.group('quote'), 'r' if raw else '', ]) # Having to do this in reverse is very irritating, but it's the # only way to make the non-greedy matches work correctly. return (RE_EXPECT.sub(replace, src[:i][::-1], count=1)[::-1] + src[i:], delta[0]) class TestCase(unittest.TestCase): longMessage = True def assertExpectedInline(self, actual, expect, skip=0): if ACCEPT: if actual != expect: # current frame and parent frame, plus any requested skip tb = traceback.extract_stack(limit=2 + skip) fn, lineno, _, _ = tb[0] print("Accepting new output for {} at {}:{}".format(self.id(), fn, lineno)) with open(fn, 'r+') as f: old = f.read() # compute the change in lineno lineno = EDIT_HISTORY.adjust_lineno(fn, lineno) new, delta = replace_string_literal(old, lineno, actual) assert old != new, "Failed to substitute string at {}:{}".format(fn, lineno) # Only write the backup file the first time we hit the # file if not EDIT_HISTORY.seen_file(fn): with open(fn + ".bak", 'w') as f_bak: f_bak.write(old) f.seek(0) f.truncate(0) f.write(new) EDIT_HISTORY.record_edit(fn, lineno, delta) else: help_text = ("To accept the new output, re-run test with " "envvar EXPECTTEST_ACCEPT=1 (we recommend " "staging/committing your changes before doing this)") if hasattr(self, "assertMultiLineEqual"): self.assertMultiLineEqual(expect, actual, msg=help_text) else: self.assertEqual(expect, actual, msg=help_text) if __name__ == "__main__": import doctest doctest.testmod()