Repository: asottile/add-trailing-comma
Pull request #78 (2019-06-08 21:57:41)

 Remove unnecessary commas on single-line things


Patch:
diff --git a/README.md b/README.md
index 09be88f..fb55078 100644
--- a/README.md
+++ b/README.md
@@ -198,3 +198,14 @@ Note that this would cause a **`SyntaxError`** in earlier python versions.
 -    ]
 +]
 ```
+
+### remove unnecessary commas
+
+yes yes, I realize the tool is called `add-trailing-comma` :laughing:
+
+```diff
+-[1, 2, 3,]
+-[1, 2, 3, ]
++[1, 2, 3]
++[1, 2, 3]
+```
diff --git a/add_trailing_comma.py b/add_trailing_comma.py
index fe36baf..96e913f 100644
--- a/add_trailing_comma.py
+++ b/add_trailing_comma.py
@@ -7,7 +7,6 @@ import ast
 import collections
 import io
 import sys
-import warnings
 
 from tokenize_rt import ESCAPED_NL
 from tokenize_rt import Offset
@@ -20,8 +19,9 @@ from tokenize_rt import UNIMPORTANT_WS
 Call = collections.namedtuple('Call', ('node', 'star_args', 'arg_offsets'))
 Func = collections.namedtuple('Func', ('node', 'star_args', 'arg_offsets'))
 Class = collections.namedtuple('Class', ('node', 'star_args', 'arg_offsets'))
-Literal = collections.namedtuple('Literal', ('node',))
-Fix = collections.namedtuple('Fix', ('braces', 'multi_arg', 'initial_indent'))
+Fix = collections.namedtuple(
+    'Fix', ('braces', 'multi_arg', 'remove_comma', 'initial_indent'),
+)
 
 NEWLINES = frozenset((ESCAPED_NL, 'NEWLINE', 'NL'))
 NON_CODING_TOKENS = frozenset(('COMMENT', ESCAPED_NL, 'NL', UNIMPORTANT_WS))
@@ -31,9 +31,7 @@ END_BRACES = frozenset((')', '}', ']'))
 
 
 def ast_parse(contents_text):
-    with warnings.catch_warnings():
-        warnings.simplefilter('ignore')
-        return ast.parse(contents_text.encode('UTF-8'))
+    return ast.parse(contents_text.encode('UTF-8'))
 
 
 def _to_offset(node):
@@ -73,7 +71,7 @@ class FindNodes(ast.NodeVisitor):
 
     def _visit_literal(self, node, key='elts'):
         if getattr(node, key):
-            self.literals[_to_offset(node)] = Literal(node)
+            self.literals[_to_offset(node)] = node
         self.generic_visit(node)
 
     visit_Set = visit_List = _visit_literal
@@ -85,9 +83,9 @@ class FindNodes(ast.NodeVisitor):
         if node.elts:
             # in < py38 tuples lie about offset -- later we must backtrack
             if sys.version_info < (3, 8):  # pragma: no cover (<py38)
-                self.tuples[_to_offset(node)] = Literal(node)
+                self.tuples[_to_offset(node)] = node
             else:  # pragma: no cover (py38+)
-                self.literals[_to_offset(node)] = Literal(node)
+                self.literals[_to_offset(node)] = node
         self.generic_visit(node)
 
     def visit_Call(self, node):
@@ -192,8 +190,18 @@ def _find_simple(first_brace, tokens):
     last_brace = i
 
     # Check if we're actually multi-line
-    if tokens[first_brace].line == tokens[last_brace].line:
-        return
+    if (
+            # we were single line, but with an extra comma and or whitespace
+            tokens[first_brace].line == tokens[last_brace].line and (
+                tokens[last_brace - 1].name == UNIMPORTANT_WS or
+                tokens[last_brace - 1].src == ','
+            )
+    ):
+        remove_comma = True
+    elif tokens[first_brace].line == tokens[last_brace].line:
+        return None
+    else:
+        remove_comma = False
 
     # determine the initial indentation
     i = first_brace
@@ -205,7 +213,12 @@ def _find_simple(first_brace, tokens):
     else:
         initial_indent = 0
 
-    return Fix((first_brace, last_brace), multi_arg, initial_indent)
+    return Fix(
+        (first_brace, last_brace),
+        multi_arg=multi_arg,
+        remove_comma=remove_comma,
+        initial_indent=initial_indent,
+    )
 
 
 def _find_call(call, i, tokens):
@@ -265,7 +278,9 @@ def _find_import(i, tokens):
         raise AssertionError('Past end?')
 
 
-def _fix_brace(fix_data, add_comma, tokens):
+def _fix_brace(tokens, fix_data, add_comma, remove_comma):
+    if fix_data is None:
+        return
     first_brace, last_brace = fix_data.braces
 
     # Figure out if either of the braces are "hugging"
@@ -279,7 +294,9 @@ def _fix_brace(fix_data, add_comma, tokens):
             tokens[last_brace - 1].src in END_BRACES or
             # Don't unhug when containing a single token (such as a triple
             # quoted string).
-            first_brace + 2 == last_brace
+            first_brace + 2 == last_brace or
+            # don't unhug if it is a single line
+            fix_data.remove_comma
     ):
         hug_open = hug_close = False
 
@@ -344,6 +361,18 @@ def _fix_brace(fix_data, add_comma, tokens):
         new_indent = fix_data.initial_indent * ' '
         tokens[last_brace - 1] = back_1._replace(src=new_indent)
 
+    if fix_data.remove_comma:
+        start = last_brace
+        if tokens[start - 1].name == UNIMPORTANT_WS:
+            start -= 1
+        if remove_comma and tokens[start - 1].src == ',':
+            start -= 1
+        del tokens[start:last_brace]
+
+
+def _one_el_tuple(node):
+    return isinstance(node, ast.Tuple) and len(node.elts) == 1
+
 
 def _changing_list(lst):
     i = 0
@@ -367,43 +396,62 @@ def _fix_src(contents_text, py35_plus, py36_plus):
         if not token.src:
             continue
 
-        fixes = []
         if token.offset in visitor.calls:
             for call in visitor.calls[token.offset]:
                 # Only fix stararg calls if asked to
                 add_comma = not call.star_args or py35_plus
-                fixes.append((add_comma, _find_call(call, i, tokens)))
+                _fix_brace(
+                    tokens, _find_call(call, i, tokens),
+                    add_comma=add_comma,
+                    remove_comma=True,
+                )
         elif token.offset in visitor.funcs:
             func = visitor.funcs[token.offset]
             add_comma = not func.star_args or py36_plus
             # functions can be treated as calls
-            fixes.append((add_comma, _find_call(func, i, tokens)))
+            _fix_brace(
+                tokens, _find_call(func, i, tokens),
+                add_comma=add_comma,
+                remove_comma=True,
+            )
         elif token.offset in visitor.classes:
             # classes can be treated as calls
             cls = visitor.classes[token.offset]
-            fixes.append((True, _find_call(cls, i, tokens)))
+            _fix_brace(
+                tokens, _find_call(cls, i, tokens),
+                add_comma=True,
+                remove_comma=True,
+            )
         elif token.offset in visitor.literals and token.src in START_BRACES:
-            fixes.append((True, _find_simple(i, tokens)))
+            _fix_brace(
+                tokens, _find_simple(i, tokens),
+                add_comma=True,
+                remove_comma=not _one_el_tuple(visitor.literals[token.offset]),
+            )
         elif token.offset in visitor.imports:
             # some imports do not have parens
-            fix = _find_import(i, tokens)
-            if fix:
-                fixes.append((True, fix))
+            _fix_brace(
+                tokens, _find_import(i, tokens),
+                add_comma=True,
+                remove_comma=True,
+            )
         # Handle parenthesized things, unhug of tuples, and comprehensions
         elif token.src in START_BRACES:
-            fixes.append((False, _find_simple(i, tokens)))
-
-        for add_comma, fix_data in fixes:
-            if fix_data is not None:
-                _fix_brace(fix_data, add_comma, tokens)
+            _fix_brace(
+                tokens, _find_simple(i, tokens),
+                add_comma=False,
+                remove_comma=False,
+            )
 
         # 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)
-            fix_data = _find_tuple(i, tokens)
-            if fix_data is not None:
-                _fix_brace(fix_data, True, tokens)
+            _fix_brace(
+                tokens, _find_tuple(i, tokens),
+                add_comma=True,
+                remove_comma=not _one_el_tuple(visitor.tuples[token.offset]),
+            )
 
     return tokens_to_src(tokens)
 


Test patch:
diff --git a/tests/add_trailing_comma_test.py b/tests/add_trailing_comma_test.py
index 39085c5..53330e7 100644
--- a/tests/add_trailing_comma_test.py
+++ b/tests/add_trailing_comma_test.py
@@ -795,6 +795,25 @@ def test_fix_classes(src, expected):
     assert _fix_src(src, py35_plus=False, py36_plus=False) == expected
 
 
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        # can't rewrite 1-element tuple
+        ('(1,)', '(1,)'),
+        # but I do want the whitespace fixed!
+        ('(1, )', '(1,)'),
+        ('(1, 2,)', '(1, 2)'),
+        ('[1, 2,]', '[1, 2]'),
+        ('[1, 2,   ]', '[1, 2]'),
+        ('{1, 2, }', '{1, 2}'),
+        ('{1: 2, }', '{1: 2}'),
+        ('f(1, 2,)', 'f(1, 2)'),
+    ),
+)
+def test_remove_extra_comma(src, expected):
+    assert _fix_src(src, py35_plus=False, py36_plus=False) == expected
+
+
 @xfailif_py2
 @pytest.mark.parametrize(
     ('src', 'expected'),