From 3b00112ac5eb8a30fc0c2630227e6337a6752491 Mon Sep 17 00:00:00 2001
From: GiGaGon <107241144+MeGaGiGaGon@users.noreply.github.com>
Date: Tue, 24 Dec 2024 12:25:08 -0800
Subject: [PATCH] Fix crash on formatting certain with statements (#4538)

Fixes #3678
---
 CHANGES.md                               |  3 +++
 src/black/linegen.py                     |  4 ++++
 src/black/nodes.py                       | 22 ++++++++++++++++++++
 tests/data/cases/remove_with_brackets.py | 26 ++++++++++++++++++++++++
 4 files changed, 55 insertions(+)

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
 
 <!-- Changes that affect Black's 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