Repository: asottile/add-trailing-comma
Pull request #89 (2019-10-29 20:00:23)

[multi-slice] Recognize multi-slice tuples
makes generics work better


Patch:
diff --git a/add_trailing_comma.py b/add_trailing_comma.py
index ebbfb61..d45038a 100644
--- a/add_trailing_comma.py
+++ b/add_trailing_comma.py
@@ -103,8 +103,10 @@ class FindNodes(ast.NodeVisitor):
     def visit_Tuple(self, node):
         # type: (ast.Tuple) -> None
         if node.elts:
+            if _to_offset(node) == _to_offset(node.elts[0]):
+                self.tuples[_to_offset(node)] = node
             # in < py38 tuples lie about offset -- later we must backtrack
-            if sys.version_info < (3, 8):  # pragma: no cover (<py38)
+            elif sys.version_info < (3, 8):  # pragma: no cover (<py38)
                 self.tuples[_to_offset(node)] = node
             else:  # pragma: no cover (py38+)
                 self.literals[_to_offset(node)] = node
@@ -283,7 +285,7 @@ def _find_call(call, i, tokens):
     return _find_simple(first_brace, tokens)
 
 
-def _find_tuple(i, tokens):  # pragma: no cover (<py38)
+def _find_tuple(i, tokens):
     # type: (int, Sequence[Token]) -> Optional[Fix]
     # tuples are evil, we need to backtrack to find the opening paren
     i -= 1
@@ -291,7 +293,7 @@ def _find_tuple(i, tokens):  # pragma: no cover (<py38)
         i -= 1
     # Sometimes tuples don't even have a paren!
     # x = 1, 2, 3
-    if tokens[i].src != '(':
+    if tokens[i].src != '(' and tokens[i].src != '[':
         return None
 
     return _find_simple(i, tokens)
@@ -484,7 +486,7 @@ def _fix_src(contents_text, py35_plus, py36_plus):
         # need to handle tuples afterwards as tuples report their starting
         # starting index as the first element, which may be one of the above
         # things.
-        if token.offset in visitor.tuples:  # pragma: no cover (<py38)
+        if token.offset in visitor.tuples:
             _fix_brace(
                 tokens, _find_tuple(i, tokens),
                 add_comma=True,


Test patch:
diff --git a/tests/add_trailing_comma_test.py b/tests/add_trailing_comma_test.py
index b1062a2..0757f8a 100644
--- a/tests/add_trailing_comma_test.py
+++ b/tests/add_trailing_comma_test.py
@@ -145,6 +145,7 @@ def test_fixes_calls(src, expected):
         'if True:\n'
         '    pass\n'
         '[x] = {y}',
+        pytest.param('x[1, 2, 3, 4]', id='multi-slice'),
     ),
 )
 def test_noop_literals(src):
@@ -204,6 +205,42 @@ def test_noop_literals(src):
             '    (),\n'
             ')',
         ),
+        pytest.param(
+            'x[\n'
+            '    1,\n'
+            '    2,\n'
+            '    3\n'
+            ']',
+            'x[\n'
+            '    1,\n'
+            '    2,\n'
+            '    3,\n'
+            ']',
+            id='multi-line multi-slice adds comma at end',
+        ),
+        pytest.param(
+            'x[1, 2, 3, ]',
+            'x[1, 2, 3]',
+            id='single line with trailing comma with space removes comma',
+        ),
+        pytest.param(
+            'x[1, 2, 3,]',
+            'x[1, 2, 3]',
+            id='single line with trailing comma with no space removes comma',
+        ),
+        pytest.param(
+            'x[\n'
+            '    (1,),\n'
+            '    2,\n'
+            '    3\n'
+            ']',
+            'x[\n'
+            '    (1,),\n'
+            '    2,\n'
+            '    3,\n'
+            ']',
+            id='nested tuple',
+        ),
     ),
 )
 def test_fixes_literals(src, expected):