Merge pull request #9 from asottile/adjust_end_brace

Adjust trailing brace indentation to match opening brace

README.md

@@ -146,9 +146,7 @@ being a syntax error.
 +)
 ```
 
-## Planned features
-
-### Match closing brace indentation
+### match closing brace indentation
 
 ```diff
  x = [

add_trailing_comma.py

@@ -310,6 +310,21 @@ def _fix_comma_and_unhug(fix_data, add_comma, tokens):
         tokens.insert(i + 1, Token('OP', ','))
 
 
+def _fix_trailing_brace(fix_data, tokens):
+    _, last_brace = fix_data.braces
+
+    back_1 = tokens[last_brace - 1]
+    back_2 = tokens[last_brace - 2]
+
+    if (
+            back_1.name == UNIMPORTANT_WS and
+            back_2.name == 'NL' and
+            len(back_1.src) != fix_data.initial_indent
+    ):
+        new_indent = fix_data.initial_indent * ' '
+        tokens[last_brace - 1] = back_1._replace(src=new_indent)
+
+
 def _fix_src(contents_text, py35_plus):
     try:
         ast_obj = ast_parse(contents_text)
@@ -343,6 +358,21 @@ def _fix_src(contents_text, py35_plus):
         if fix_data is not None:
             _fix_comma_and_unhug(fix_data, add_comma, tokens)
 
+    # Need a second pass to fix trailing braces after indentation is fixed
+    for i, token in reversed(tuple(enumerate(tokens))):
+        key = Offset(token.line, token.utf8_byte_offset)
+        fix_data = None
+
+        if key in visitor.calls:
+            fix_data = _find_call(visitor.calls[key], i, tokens)
+        elif key in visitor.funcs:
+            fix_data = _find_call(visitor.funcs[key], i, tokens)
+        elif key in visitor.literals:
+            fix_data = _find_literal(visitor.literals[key], i, tokens)
+
+        if fix_data is not None:
+            _fix_trailing_brace(fix_data, tokens)
+
     return tokens_to_src(tokens)
 
 

tests/add_trailing_comma_test.py

@@ -393,14 +393,13 @@ def test_noop_unhugs(src):
             ' },\n'
             ' }',
 
-            # TODO: need to adjust trailing braces
             '{\n'
             "    'foo': 'bar',\n"
             "    'baz':\n"
             '       {\n'
             "          'id': 1,\n"
-            '    },\n'
-            '    }',
+            '       },\n'
+            '}',
         ),
         (
             'f(g(\n'
@@ -466,6 +465,42 @@ def test_fix_unhugs_py3_only(src, expected):
     assert _fix_src(src, py35_plus=False) == expected
 
 
+@pytest.mark.parametrize(
+    'src',
+    (
+        '[]',
+        'x = [\n'
+        '    1, 2, 3,\n'
+        ']',
+        'y = [\n'
+        '    [\n'
+        '        1, 2, 3, 4,\n'
+        '    ],\n'
+        ']',
+    ),
+)
+def test_noop_trailing_brace(src):
+    assert _fix_src(src, py35_plus=False) == src
+
+
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        (
+            'x = [\n'
+            '    1,\n'
+            '    ]',
+
+            'x = [\n'
+            '    1,\n'
+            ']',
+        ),
+    ),
+)
+def test_fix_trailing_brace(src, expected):
+    assert _fix_src(src, py35_plus=False) == expected
+
+
 def test_main_trivial():
     assert main(()) == 0