diff --git a/CHANGES.md b/CHANGES.md index c1e6320..d2955d2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -13,6 +13,9 @@ - Fix formatting cells in IPython notebooks with magic methods and starting or trailing empty lines (#4484) +- Fix crash when formatting `with` statements containing tuple generators/unpacking + (#4538) + ### Preview style diff --git a/src/black/linegen.py b/src/black/linegen.py index 7bd018d..6bd2e5c 100644 --- a/src/black/linegen.py +++ b/src/black/linegen.py @@ -45,6 +45,7 @@ is_atom_with_invisible_parens, is_docstring, is_empty_tuple, + is_generator, is_lpar_token, is_multiline_string, is_name_token, @@ -55,6 +56,7 @@ is_rpar_token, is_stub_body, is_stub_suite, + is_tuple_containing_star, is_tuple_containing_walrus, is_type_ignore_comment_string, is_vararg, @@ -1628,6 +1630,8 @@ def maybe_make_parens_invisible_in_atom( and max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY ) or is_tuple_containing_walrus(node) + or is_tuple_containing_star(node) + or is_generator(node) ): return False diff --git a/src/black/nodes.py b/src/black/nodes.py index 927b9ee..f3688cd 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -621,6 +621,28 @@ def is_tuple_containing_walrus(node: LN) -> bool: return any(child.type == syms.namedexpr_test for child in gexp.children) +def is_tuple_containing_star(node: LN) -> bool: + """Return True if `node` holds a tuple that contains a star operator.""" + if node.type != syms.atom: + return False + gexp = unwrap_singleton_parenthesis(node) + if gexp is None or gexp.type != syms.testlist_gexp: + return False + + return any(child.type == syms.star_expr for child in gexp.children) + + +def is_generator(node: LN) -> bool: + """Return True if `node` holds a generator.""" + if node.type != syms.atom: + return False + gexp = unwrap_singleton_parenthesis(node) + if gexp is None or gexp.type != syms.testlist_gexp: + return False + + return any(child.type == syms.old_comp_for for child in gexp.children) + + def is_one_sequence_between( opening: Leaf, closing: Leaf, diff --git a/tests/data/cases/remove_with_brackets.py b/tests/data/cases/remove_with_brackets.py index ea58ab9..f2319e0 100644 --- a/tests/data/cases/remove_with_brackets.py +++ b/tests/data/cases/remove_with_brackets.py @@ -53,6 +53,19 @@ with ((((CtxManager1()))) as example1, (((CtxManager2()))) as example2): ... +# regression tests for #3678 +with (a, *b): + pass + +with (a, (b, *c)): + pass + +with (a for b in c): + pass + +with (a, (b for c in d)): + pass + # output with open("bla.txt"): pass @@ -117,3 +130,16 @@ with CtxManager1() as example1, CtxManager2() as example2: ... + +# regression tests for #3678 +with (a, *b): + pass + +with a, (b, *c): + pass + +with (a for b in c): + pass + +with a, (b for c in d): + pass