Merge pull request #53 from asottile/from_import_trailing_comma

Add trailing commas for from imports

README.md

@@ -148,6 +148,16 @@ following change:
 
 Note that this would cause a **`SyntaxError`** in earlier python versions.
 
+### trailing commas for `from` imports
+
+```diff
+ from os import (
+     path,
+-    makedirs
++    makedirs,
+ )
+```
+
 ### unhug trailing paren
 
 ```diff

add_trailing_comma.py

@@ -64,6 +64,7 @@ class FindNodes(ast.NodeVisitor):
         self.funcs = {}
         self.literals = {}
         self.tuples = {}
+        self.imports = set()
 
     def _visit_literal(self, node, key='elts'):
         if getattr(node, key):
@@ -143,6 +144,10 @@ class FindNodes(ast.NodeVisitor):
 
         self.generic_visit(node)
 
+    def visit_ImportFrom(self, node):
+        self.imports.add(Offset(node.lineno, node.col_offset))
+        self.generic_visit(node)
+
 
 def _find_simple(first_brace, tokens):
     brace_stack = [first_brace]
@@ -227,6 +232,18 @@ def _find_tuple(i, tokens):
     return _find_simple(i, tokens)
 
 
+def _find_import(i, tokens):
+    # progress forwards until we find either a `(` or a newline
+    for i in range(i, len(tokens)):
+        token = tokens[i]
+        if token.name == 'NEWLINE':
+            return
+        elif token.name == 'OP' and token.src == '(':
+            return _find_simple(i, tokens)
+    else:
+        raise AssertionError('Past end?')
+
+
 def _fix_brace(fix_data, add_comma, tokens):
     first_brace, last_brace = fix_data.braces
 
@@ -346,6 +363,11 @@ def _fix_src(contents_text, py35_plus, py36_plus):
         # Handle parenthesized things, unhug of tuples, and comprehensions
         elif token.src in START_BRACES:
             fixes.append((False, _find_simple(i, tokens)))
+        elif key in visitor.imports:
+            # some imports do not have parens
+            fix = _find_import(i, tokens)
+            if fix:
+                fixes.append((True, fix))
 
         for add_comma, fix_data in fixes:
             if fix_data is not None:

tests/add_trailing_comma_test.py

@@ -684,6 +684,72 @@ def test_fix_trailing_brace(src, expected):
     assert _fix_src(src, py35_plus=False, py36_plus=False) == expected
 
 
+@pytest.mark.parametrize(
+    'src',
+    (
+        'from os import path, makedirs\n',
+        'from os import (path, makedirs)\n',
+        'from os import (\n'
+        '    path,\n'
+        '    makedirs,\n'
+        ')',
+    ),
+)
+def test_fix_from_import_noop(src):
+    assert _fix_src(src, py35_plus=False, py36_plus=False) == src
+
+
+@pytest.mark.parametrize(
+    ('src', 'expected'),
+    (
+        (
+            'from os import (\n'
+            '    makedirs,\n'
+            '    path\n'
+            ')',
+            'from os import (\n'
+            '    makedirs,\n'
+            '    path,\n'
+            ')',
+        ),
+        (
+            'from os import \\\n'
+            '   (\n'
+            '        path,\n'
+            '        makedirs\n'
+            '   )\n',
+            'from os import \\\n'
+            '   (\n'
+            '        path,\n'
+            '        makedirs,\n'
+            '   )\n',
+        ),
+        (
+            'from os import (\n'
+            '    makedirs,\n'
+            '    path,\n'
+            '    )',
+            'from os import (\n'
+            '    makedirs,\n'
+            '    path,\n'
+            ')',
+        ),
+        (
+            'if True:\n'
+            '    from os import (\n'
+            '        makedirs\n'
+            '    )',
+            'if True:\n'
+            '    from os import (\n'
+            '        makedirs,\n'
+            '    )',
+        ),
+    ),
+)
+def test_fix_from_import(src, expected):
+    assert _fix_src(src, py35_plus=False, py36_plus=False) == expected
+
+
 def test_main_trivial():
     assert main(()) == 0