Repository: asottile/add-trailing-comma
Pull request #6 (2017-07-13 01:18:35)

Unhug leading and trailing braces
CC @chriskuehl 

This makes a best effort to reindent in the opening hug case.  There's some oddities that aren't perfect, (they're pretty unlikely and hopefully will get noticed / fixed by some other tool such as autopep8).

I ran it against pre-commit/pre-commit and this was the diff generated (vs output against 0.3.0) (without involving autopep8):

```diff
--- a/tests/commands/run_test.py
+++ b/tests/commands/run_test.py
@@ -86,8 +86,10 @@ def _do_run(cap_out, repo, args, environ={}, config_file=C.CONFIG_FILE):
     return ret, printed
 
 
-def _test_run(cap_out, repo, opts, expected_outputs, expected_ret, stage,
-              config_file=C.CONFIG_FILE):
+def _test_run(
+    cap_out, repo, opts, expected_outputs, expected_ret, stage,
+    config_file=C.CONFIG_FILE,
+):
     if stage:
         stage_a_file()
     args = _get_opts(**opts)
@@ -571,8 +573,10 @@ def test_lots_of_files(mock_out_store_directory, tempdir_factory):
 
 
 @pytest.mark.parametrize(
-    ('hook_stage', 'stage_for_first_hook', 'stage_for_second_hook',
-     'expected_output'),
+    (
+        'hook_stage', 'stage_for_first_hook', 'stage_for_second_hook',
+        'expected_output',
+    ),
     (
         ('push', ['commit'], ['commit'], [b'', b'']),
         ('push', ['commit', 'push'], ['commit', 'push'],
```

suntzu86: +1
Can't sendlove anymore but this is super awesome!
asottile: It now better handles reindenting by considering the *least* indented indent inside the call/literal to be the baseline and then adjusting the indent from there.  One last thing I'd like to do is fix the trailing brace:

```diff
 x = [
     1,
     2,
     3,
-    ]
+]
```

Patch:
diff --git a/README.md b/README.md
index 56316a4..0160d8a 100644
--- a/README.md
+++ b/README.md
@@ -125,9 +125,6 @@ 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
@@ -148,3 +145,16 @@ being a syntax error.
 +    arg2,
 +)
 ```
+
+## Planned features
+
+### Match closing brace indentation
+
+```diff
+ x = [
+     1,
+     2,
+     3,
+-    ]
++]
+```
diff --git a/add_trailing_comma.py b/add_trailing_comma.py
index 26d7f83..49e98c7 100644
--- a/add_trailing_comma.py
+++ b/add_trailing_comma.py
@@ -19,7 +19,10 @@ Func = collections.namedtuple('Func', ('node', 'arg_offsets'))
 Literal = collections.namedtuple('Literal', ('node', 'braces', 'backtrack'))
 Literal.__new__.__defaults__ = (False,)
 
+NEWLINES = frozenset(('NEWLINE', 'NL'))
 NON_CODING_TOKENS = frozenset(('COMMENT', 'NL', UNIMPORTANT_WS))
+INDENT_TOKENS = frozenset(('INDENT', UNIMPORTANT_WS))
+END_BRACES = frozenset((')', '}', ']'))
 
 
 def ast_parse(contents_text):
@@ -159,10 +162,10 @@ class FindNodes(ast.NodeVisitor):
         self.generic_visit(node)
 
 
-def _fix_inner(brace_start, brace_end, first_paren, tokens):
-    brace_stack = [first_paren]
+def _fix_inner(brace_start, brace_end, first_brace, tokens):
+    brace_stack = [first_brace]
 
-    for i in range(first_paren + 1, len(tokens)):
+    for i in range(first_brace + 1, len(tokens)):
         token = tokens[i]
         if token.src == brace_start:
             brace_stack.append(i)
@@ -174,17 +177,69 @@ def _fix_inner(brace_start, brace_end, first_paren, tokens):
     else:
         raise AssertionError('Past end?')
 
+    last_brace = i
+
     # This was not actually a multi-line call, despite the ast telling us that
-    if tokens[first_paren].line == tokens[i].line:
+    if tokens[first_brace].line == tokens[last_brace].line:
         return
 
+    # Figure out if either of the braces are "hugging"
+    hug_open = tokens[first_brace + 1].name not in NON_CODING_TOKENS
+    hug_close = tokens[last_brace - 1].name not in NON_CODING_TOKENS
+    if hug_open and tokens[last_brace - 1].src in END_BRACES:
+        hug_open = hug_close = False
+
+    # determine the initial indentation
+    i = first_brace
+    while i >= 0 and tokens[i].name not in NEWLINES:
+        i -= 1
+
+    if i >= 0 and tokens[i + 1].name in INDENT_TOKENS:
+        initial_indent = len(tokens[i + 1].src)
+    else:
+        initial_indent = 0
+
+    # fix open hugging
+    if hug_open:
+        new_indent = initial_indent + 4
+
+        tokens[first_brace + 1:first_brace + 1] = [
+            Token('NL', '\n'), Token(UNIMPORTANT_WS, ' ' * new_indent),
+        ]
+        last_brace += 2
+
+        # Adust indentation for the rest of the things
+        min_indent = None
+        indents = []
+        for i in range(first_brace + 3, last_brace):
+            if tokens[i - 1].name == 'NL' and tokens[i].name == UNIMPORTANT_WS:
+                if min_indent is None:
+                    min_indent = len(tokens[i].src)
+                elif len(tokens[i].src) < min_indent:
+                    min_indent = len(tokens[i].src)
+
+                indents.append(i)
+
+        for i in indents:
+            oldlen = len(tokens[i].src)
+            newlen = oldlen - min_indent + new_indent
+            tokens[i] = tokens[i]._replace(src=' ' * newlen)
+
+    # fix close hugging
+    if hug_close:
+        tokens[last_brace:last_brace] = [
+            Token('NL', '\n'),
+            Token(UNIMPORTANT_WS, ' ' * initial_indent),
+        ]
+        last_brace += 2
+
     # From there, we can walk backwards and decide whether a comma is needed
-    i -= 1
+    i = last_brace - 1
     while tokens[i].name in NON_CODING_TOKENS:
         i -= 1
 
     # If we're not a hugging paren, we can insert a comma
-    if tokens[i].src != ',' and tokens[i + 1].src != brace_end:
+    if tokens[i].src != ',' and i + 1 != last_brace:
         tokens.insert(i + 1, Token('OP', ','))
 
 
@@ -201,7 +256,7 @@ def _fix_call(call, i, tokens):
     #     func_name(arg, arg, arg)
     #              ^ outer paren
     brace_start, brace_end = '(', ')'
-    first_paren = None
+    first_brace = None
     paren_stack = []
     for i in range(i, len(tokens)):
         token = tokens[i]
@@ -213,12 +268,12 @@ def _fix_call(call, i, tokens):
             paren_stack.pop()
 
         if (token.line, token.utf8_byte_offset) in call.arg_offsets:
-            first_paren = paren_stack[0]
+            first_brace = paren_stack[0]
             break
     else:
         raise AssertionError('Past end?')
 
-    _fix_inner(brace_start, brace_end, first_paren, tokens)
+    _fix_inner(brace_start, brace_end, first_brace, tokens)
 
 
 def _fix_literal(literal, i, tokens):


Test patch:
diff --git a/tests/add_trailing_comma_test.py b/tests/add_trailing_comma_test.py
index c414166..810b741 100644
--- a/tests/add_trailing_comma_test.py
+++ b/tests/add_trailing_comma_test.py
@@ -297,6 +297,130 @@ def test_fixes_defs(src, expected):
     assert _fix_src(src, py35_plus=False) == expected
 
 
+@pytest.mark.parametrize(
+    'src',
+    (
+        'f(x, y, z)',
+        'f(\n'
+        '    x,\n'
+        ')',
+        # Single argument, don't unhug
+        'f((\n'
+        '    1, 2, 3,\n'
+        '))',
+        'f([\n'
+        '    1, 2, 3,\n'
+        '])',
+    ),
+)
+def test_noop_unhugs(src):
+    assert _fix_src(src, py35_plus=False) == src
+
+
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        (
+            'f(\n'
+            '    a)',
+
+            'f(\n'
+            '    a,\n'
+            ')',
+        ),
+        (
+            'f(a,\n'
+            '  b,\n'
+            ')',
+
+            'f(\n'
+            '    a,\n'
+            '    b,\n'
+            ')',
+        ),
+        (
+            'f(a,\n'
+            '  b,\n'
+            '  c)',
+
+            'f(\n'
+            '    a,\n'
+            '    b,\n'
+            '    c,\n'
+            ')',
+        ),
+        # if there's already a trailing comma, don't add a new one
+        (
+            'f(\n'
+            '    a,)',
+
+            'f(\n'
+            '    a,\n'
+            ')',
+        ),
+        (
+            'with a(\n'
+            '    b,\n'
+            '    c):\n'
+            '    pass',
+
+            'with a(\n'
+            '    b,\n'
+            '    c,\n'
+            '):\n'
+            '    pass',
+        ),
+        (
+            'if True:\n'
+            '    with a(\n'
+            '        b,\n'
+            '        c):\n'
+            '        pass',
+
+            'if True:\n'
+            '    with a(\n'
+            '        b,\n'
+            '        c,\n'
+            '    ):\n'
+            '        pass',
+        ),
+        (
+            "{'foo': 'bar',\n"
+            " 'baz':\n"
+            '    {\n'
+            "       'id': 1,\n"
+            ' },\n'
+            ' }',
+
+            # TODO: need to adjust trailing braces
+            '{\n'
+            "    'foo': 'bar',\n"
+            "    'baz':\n"
+            '       {\n'
+            "          'id': 1,\n"
+            '    },\n'
+            '    }',
+        ),
+        (
+            'f(g(\n'
+            '      a,\n'
+            '  ),\n'
+            '  1,\n'
+            ')',
+
+            'f(\n'
+            '    g(\n'
+            '        a,\n'
+            '    ),\n'
+            '    1,\n'
+            ')',
+        ),
+    ),
+)
+def test_fix_unhugs(src, expected):
+    assert _fix_src(src, py35_plus=False) == expected
+
+
 def test_main_trivial():
     assert main(()) == 0