Skip to content

Commit

Permalink
reimplement quote selection logic
Browse files Browse the repository at this point in the history
  • Loading branch information
sunmy2019 committed Sep 3, 2023
1 parent 43d177c commit 21229ce
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 23 deletions.
55 changes: 32 additions & 23 deletions Lib/ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -1225,40 +1225,49 @@ def _write_str_avoiding_backslashes(self, string, *, quote_types=_ALL_QUOTES):

def visit_JoinedStr(self, node):
self.write("f")
if self._avoid_backslashes:
with self.buffered() as buffer:
self._write_fstring_inner(node)
return self._write_str_avoiding_backslashes("".join(buffer))

# If we don't need to avoid backslashes globally (i.e., we only need
# to avoid them inside FormattedValues), it's cosmetically preferred
# to use escaped whitespace. That is, it's preferred to use backslashes
# for cases like: f"{x}\n". To accomplish this, we keep track of what
# in our buffer corresponds to FormattedValues and what corresponds to
# Constant parts of the f-string, and allow escapes accordingly.

fstring_parts = []
for value in node.values:
with self.buffered() as buffer:
self._write_fstring_inner(value)
fstring_parts.append(
("".join(buffer), isinstance(value, Constant))
)
fstring_parts.append(("".join(buffer), isinstance(value, Constant)))

# We decide if we need to write a multi-line `f-string` since it is only
# necessary when we have "\n" inside formatted values.
use_multiline = any(
"\n" in value for value, is_constant in fstring_parts if not is_constant
)

# We then choose the quote type we use. We let `repr` do this work for
# now. This can be easily modified afterwards.
quote = repr(
"".join(value for value, is_constant in fstring_parts if is_constant)
)[0]
quote_type = quote * 3 if use_multiline else quote

new_fstring_parts = []
quote_types = list(_ALL_QUOTES)
for value, is_constant in fstring_parts:
if is_constant:
value, quote_types = self._str_literal_helper(
value,
quote_types=quote_types,
escape_special_whitespace=True,
)
elif "\n" in value:
quote_types = [q for q in quote_types if q in _MULTI_QUOTES]
consecutive_quotes = 0
res = []
for c in value:
if c == "\\" or not c.isprintable():
res.append(c.encode("unicode_escape").decode("ascii"))
continue
if c == quote:
if consecutive_quotes == len(quote_type) - 1:
# escape when we see a full `quote_type`
res.append("\\")
consecutive_quotes = 0
else:
consecutive_quotes += 1
else:
consecutive_quotes = 0
res.append(c)
value = "".join(res)
new_fstring_parts.append(value)

value = "".join(new_fstring_parts)
quote_type = quote_types[0]
self.write(f"{quote_type}{value}{quote_type}")

def _write_fstring_inner(self, node):
Expand Down
6 changes: 6 additions & 0 deletions Lib/test/test_unparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,10 +196,14 @@ def test_fstrings_complicated(self):
self.check_ast_roundtrip('f"""{g(\'\'\'\n\'\'\')}"""')
self.check_ast_roundtrip('''f"a\\r\\nb"''')
self.check_ast_roundtrip('''f"\\u2028{'x'}"''')
self.check_ast_roundtrip("f\"'''{1}\\\"\\\"\\\"\"")
self.check_ast_roundtrip('f\'\\\'\\\'\\\'{1}"""\'')
self.check_ast_roundtrip('f\'\'\'\'\'\\\'\'\'\\\'{x:\n}""""\'\'\'')

def test_fstrings_pep701(self):
self.check_ast_roundtrip('f" something { my_dict["key"] } something else "')
self.check_ast_roundtrip('f"{f"{f"{f"{f"{f"{1+1}"}"}"}"}"}"')
self.check_ast_roundtrip("f'{f'{f'{f'{f'{f'{1+1}'}'}'}'}'}'")

def test_strings(self):
self.check_ast_roundtrip("u'foo'")
Expand Down Expand Up @@ -518,6 +522,8 @@ def test_fstrings(self):
self.check_src_roundtrip(r"f'{x}\n'")
self.check_src_roundtrip("f'{'\\n'}\\n'")
self.check_src_roundtrip("f'{f'{x}\\n'}\\n'")
self.check_src_roundtrip('f\'\\\'\\\'\\\'{1}"""\'')
self.check_src_roundtrip('f\'\'\'\'\'\\\'\'\'\\\'{x:\n}""""\'\'\'')

def test_docstrings(self):
docstrings = (
Expand Down

0 comments on commit 21229ce

Please sign in to comment.