Repository: asottile/add-trailing-comma (language: python)

Issue:
--py39-plus: multi-with-statements
Add support for trailing commas for this:

```diff
 with (
     foo as bar,
-    baz as womp
+    baz as womp,
 ):
     # ...
```


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index bb6195f..b32fb42 100644
--- a/README.md
+++ b/README.md
@@ -178,6 +178,17 @@ Note that this would cause a **`SyntaxError`** in earlier python versions.
      pass
 ```
 
+### trailing comma for with statement
+
+```diff
+ with (
+         open('f1', 'r') as f1,
+-        open('f2', 'w') as f2
++        open('f2', 'w') as f2,
+ ):
+     pass
+```
+
 ### trailing comma for match statement
 
 ```diff
diff --git a/add_trailing_comma/_plugins/_with.py b/add_trailing_comma/_plugins/_with.py
new file mode 100644
index 0000000..ffbaf81
--- /dev/null
+++ b/add_trailing_comma/_plugins/_with.py
@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import ast
+import sys
+from typing import Iterable
+
+from tokenize_rt import Offset
+from tokenize_rt import Token
+
+from add_trailing_comma._ast_helpers import ast_to_offset
+from add_trailing_comma._data import register
+from add_trailing_comma._data import State
+from add_trailing_comma._data import TokenFunc
+from add_trailing_comma._token_helpers import find_simple
+from add_trailing_comma._token_helpers import fix_brace
+
+
+if sys.version_info >= (3, 9):  # pragma: >=3.9 cover
+    def _fix_with(i: int, tokens: list[Token]) -> None:
+        i += 1
+        if tokens[i].name == 'UNIMPORTANT_WS':
+            i += 1
+        if tokens[i].src == '(':
+            fix = find_simple(i, tokens)
+            # only fix if outer parens are for the with items (next is ':')
+            if fix is not None and tokens[fix.braces[-1] + 1].src == ':':
+                fix_brace(tokens, fix, add_comma=True, remove_comma=True)
+
+    @register(ast.With)
+    def visit_With(
+        state: State,
+        node: ast.With,
+    ) -> Iterable[tuple[Offset, TokenFunc]]:
+        yield ast_to_offset(node), _fix_with


Test patch:
diff --git a/tests/features/with_test.py b/tests/features/with_test.py
new file mode 100644
index 0000000..2f64fa2
--- /dev/null
+++ b/tests/features/with_test.py
@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+from add_trailing_comma._main import _fix_src
+
+
+@pytest.mark.parametrize(
+    'src',
+    (
+        pytest.param(
+            'from threading import Lock\n'
+            'with (Lock() as l):\n'
+            '    pass',
+            id='simple named context manager',
+        ),
+        pytest.param(
+            'with (\n'
+            '    open("wat")\n'
+            ') as f, open("2") as f2: pass',
+            id='parenthesized expression',
+        ),
+        pytest.param(
+            'with open("/tmp/t.py") as f: pass',
+            id='old style',
+        ),
+        pytest.param(
+            'with open("/tmp/t.py") as f, \\\n'
+            '     open("/tmp/y.py") as g: pass',
+            id='escaped newline',
+        ),
+        pytest.param(
+            'with (open("/tmp/t.py") as f): pass',
+            id='single item',
+        ),
+        pytest.param(
+            'with (open("/tmp/t.py") as t, open("/tmp/y.py") as y): pass',
+            id='single line',
+        ),
+    ),
+)
+def test_noop(src):
+    assert _fix_src(src, min_version=(2, 7)) == src
+
+
+@pytest.mark.xfail(sys.version_info < (3, 9), reason='py39+')
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        pytest.param(
+            # Make sure that whitespace is not expected after "with"
+            'with(\n'
+            '        open("/tmp/t.txt") as file1,\n'
+            '        open("/tmp/t.txt") as file2\n'
+            '): pass',
+
+            'with(\n'
+            '        open("/tmp/t.txt") as file1,\n'
+            '        open("/tmp/t.txt") as file2,\n'
+            '): pass',
+            id='simple usecase',
+        ),
+        pytest.param(
+            'from threading import lock\n'
+            'with (lock() as l,\n'
+            '     open("/tmp/t.txt")):\n'
+            '    pass',
+
+            'from threading import lock\n'
+            'with (\n'
+            '    lock() as l,\n'
+            '    open("/tmp/t.txt"),\n'
+            '):\n'
+            '    pass',
+            id='unhug',
+        ),
+        pytest.param(
+            'with (open(\n'
+            '    "a",\n'
+            '    some_other_really_long_parameter=True,\n'
+            ') as a, a.lock): pass',
+
+            'with (\n'
+            '    open(\n'
+            '        "a",\n'
+            '        some_other_really_long_parameter=True,\n'
+            '    ) as a, a.lock,\n'
+            '): pass',
+            id='lower level linebreaks',
+        ),
+        pytest.param(
+            'with (a as b, c as d,): pass\n',
+            'with (a as b, c as d): pass\n',
+            id='remove unnecessary comma',
+        ),
+        pytest.param(
+            'with (a as b,): pass\n',
+            'with (a as b): pass\n',
+            id='remove unnecessary comma one item',
+        ),
+    ),
+)
+def test_py39_multiwith(src, expected):
+    assert _fix_src(src, min_version=(2, 7)) == expected