Repository: asottile/add-trailing-comma
Pull request #5 (2017-07-12 04:09:07)

Also add trailing commas for FunctionDef


Patch:
diff --git a/README.md b/README.md
index 886f708..1f290b4 100644
--- a/README.md
+++ b/README.md
@@ -111,9 +111,6 @@ Note that this would cause a **`SyntaxError`** in earlier python versions.
  ]
 ```
 
-
-## Planned features
-
 ### trailing commas for function definitions
 
 ```diff
@@ -124,6 +121,13 @@ Note that this would cause a **`SyntaxError`** in earlier python versions.
  ):
 ```
 
+Note that functions with starargs (`*args`), kwargs (`**kwargs`), or python 3
+keyword-only arguments (`..., *, ...`) cannot have a trailing comma due to it
+being a syntax error.
+
+
+## Planned features
+
 ### unhug trailing paren
 
 ```diff
diff --git a/add_trailing_comma.py b/add_trailing_comma.py
index 14373be..34c1dc7 100644
--- a/add_trailing_comma.py
+++ b/add_trailing_comma.py
@@ -15,6 +15,7 @@ from tokenize_rt import UNIMPORTANT_WS
 
 Offset = collections.namedtuple('Offset', ('line', 'utf8_byte_offset'))
 Call = collections.namedtuple('Call', ('node', 'star_args', 'arg_offsets'))
+Func = collections.namedtuple('Func', ('node', 'arg_offsets'))
 Literal = collections.namedtuple('Literal', ('node', 'braces', 'backtrack'))
 Literal.__new__.__defaults__ = (False,)
 
@@ -53,6 +54,7 @@ def _is_star_star_kwarg(node):
 class FindNodes(ast.NodeVisitor):
     def __init__(self):
         self.calls = {}
+        self.funcs = {}
         self.literals = {}
         self.has_new_syntax = False
 
@@ -135,6 +137,27 @@ class FindNodes(ast.NodeVisitor):
 
         self.generic_visit(node)
 
+    def visit_FunctionDef(self, node):
+        has_starargs = (
+            node.args.vararg or node.args.kwarg or
+            # python 3 only
+            getattr(node.args, 'kwonlyargs', None)
+        )
+        orig = node.lineno
+        is_multiline = False
+        offsets = set()
+        for argnode in node.args.args:
+            offset = _to_offset(argnode)
+            if offset.line > orig:
+                is_multiline = True
+            offsets.add(offset)
+
+        if is_multiline and not has_starargs:
+            key = Offset(node.lineno, node.col_offset)
+            self.funcs[key] = Func(node, offsets)
+
+        self.generic_visit(node)
+
 
 def _fix_inner(brace_start, brace_end, first_paren, tokens):
     i = first_paren
@@ -235,6 +258,9 @@ def _fix_commas(contents_text, py35_plus):
                 _fix_call(call, i, tokens)
         elif key in visitor.literals:
             _fix_literal(visitor.literals[key], i, tokens)
+        elif key in visitor.funcs:
+            # functions can be treated as calls
+            _fix_call(visitor.funcs[key], i, tokens)
 
     return tokens_to_src(tokens)
 


Test patch:
diff --git a/tests/add_trailing_comma_test.py b/tests/add_trailing_comma_test.py
index 5d1c798..739247e 100644
--- a/tests/add_trailing_comma_test.py
+++ b/tests/add_trailing_comma_test.py
@@ -253,6 +253,50 @@ def test_noop_tuple_literal_without_braces():
     assert _fix_commas(src, py35_plus=False) == src
 
 
+@pytest.mark.parametrize(
+    'src',
+    (
+        'def f(arg1, arg2): pass',
+        'def f(\n'
+        '        arg1,\n'
+        '        arg2,\n'
+        '): pass',
+        # *args forbid trailing commas
+        'def f(\n'
+        '        *args\n'
+        '): pass'
+        # **kwargs forbid trailing commas
+        'def f(\n'
+        '        **kwargs\n'
+        '): pass',
+        # keyword-only args forbid trailing commas (or are py2 syntax error)
+        'def f(\n'
+        '        *, arg=1\n'
+        '): pass',
+    ),
+)
+def test_noop_function_defs(src):
+    assert _fix_commas(src, py35_plus=False) == src
+
+
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        (
+            'def f(\n'
+            '        x\n'
+            '): pass',
+
+            'def f(\n'
+            '        x,\n'
+            '): pass',
+        ),
+    ),
+)
+def test_fixes_defs(src, expected):
+    assert _fix_commas(src, py35_plus=False) == expected
+
+
 def test_main_trivial():
     assert main(()) == 0