Merge pull request #232 from timgraham/221

Fix #221 -- Add include_trailing_comma setting.

5 files changed, 87 insertions, 4 deletions

diff --git a/README.md b/README.md
index 47358cc6..f055af35 100644
--- a/README.md
+++ b/README.md
@@ -263,6 +263,10 @@ For example:
 
 is equivalent to 4
 
+For the import styles that use parentheses, you can control whether or not to
+include a trailing comma after the last import with the include_trailing_comma
+option (defaults to false).
+
 Intelligently Balanced Multi-line Imports
 ======================
 
diff --git a/isort/isort.py b/isort/isort.py
index e9aa9e91..beb36bc5 100644
--- a/isort/isort.py
+++ b/isort/isort.py
@@ -477,11 +477,16 @@ class SortImports(object):
                 comments = None
             else:
                 statement += ", " + next_import
-        return statement + ")"
+        return statement + ("," if self.config['include_trailing_comma'] else "") + ")"
 
     def _output_vertical(self, statement, imports, white_space, indent, line_length, comments):
         first_import = self._add_comments(comments, imports.pop(0) + ",") + "\n" + white_space
-        return "{0}({1}{2})".format(statement, first_import, (",\n" + white_space).join(imports))
+        return "{0}({1}{2}{3})".format(
+            statement,
+            first_import,
+            (",\n" + white_space).join(imports),
+            "," if self.config['include_trailing_comma'] else "",
+        )
 
     def _output_hanging_indent(self, statement, imports, white_space, indent, line_length, comments):
         statement += imports.pop(0)
@@ -496,8 +501,13 @@ class SortImports(object):
         return statement
 
     def _output_vertical_hanging_indent(self, statement, imports, white_space, indent, line_length, comments):
-        return "{0}({1}\n{2}{3}\n)".format(statement, self._add_comments(comments), indent,
-                                           (",\n" + indent).join(imports))
+        return "{0}({1}\n{2}{3}{4}\n)".format(
+            statement,
+            self._add_comments(comments),
+            indent,
+            (",\n" + indent).join(imports),
+            "," if self.config['include_trailing_comma'] else "",
+         )
 
     def _output_vertical_grid_common(self, statement, imports, white_space, indent, line_length, comments):
         statement += self._add_comments(comments, "(") + "\n" + indent + imports.pop(0)
@@ -507,6 +517,8 @@ class SortImports(object):
             if len(next_statement.split("\n")[-1]) + 1 > line_length:
                 next_statement = "{0},\n{1}{2}".format(statement, indent, next_import)
             statement = next_statement
+        if self.config['include_trailing_comma']:
+            statement += ','
         return statement
 
     def _output_vertical_grid(self, statement, imports, white_space, indent, line_length, comments):
diff --git a/isort/main.py b/isort/main.py
index 2d1a92f8..727f01a0 100755
--- a/isort/main.py
+++ b/isort/main.py
@@ -99,6 +99,8 @@ def main():
                         help="Ensures the output doesn't save if the resulting file contains syntax errors.")
     parser.add_argument('-cs', '--combine-star', dest='combine_star', action='store_true',
                         help="Ensures that if a star import is present, nothing else is imported from that namespace.")
+    parser.add_argument('-tc', '--trailing-comma', dest='trailing_comma', action='store_true',
+                        help='Includes a trailing comma on multi line imports that include parentheses.')
     parser.add_argument('-v', '--version', action='version', version='isort {0}'.format(__version__))
     parser.add_argument('-vb', '--verbose', action='store_true', dest="verbose",
                         help='Shows verbose output, such as when files are skipped or when a check is successful.')
diff --git a/isort/settings.py b/isort/settings.py
index f3aa3ae3..664267ba 100644
--- a/isort/settings.py
+++ b/isort/settings.py
@@ -86,6 +86,7 @@ default = {'force_to_top': [],
            'lines_after_imports': -1,
            'combine_as_imports': False,
            'combine_star': False,
+           'include_trailing_comma': False,
            'verbose': False}
 
 
diff --git a/test_isort.py b/test_isort.py
index 37fd6dc3..180c506f 100644
--- a/test_isort.py
+++ b/test_isort.py
@@ -27,6 +27,7 @@ from pies.overrides import *
 from isort.isort import SortImports
 from isort.settings import WrapModes
 
+SHORT_IMPORT = "from third_party import lib1, lib2, lib3, lib4"
 REALLY_LONG_IMPORT = ("from third_party import lib1, lib2, lib3, lib4, lib5, lib6, lib7, lib8, lib9, lib10, lib11,"
                       "lib12, lib13, lib14, lib15, lib16, lib17, lib18, lib20, lib21, lib22")
 REALLY_LONG_IMPORT_WITH_COMMENT = ("from third_party import lib1, lib2, lib3, lib4, lib5, lib6, lib7, lib8, lib9, "
@@ -937,6 +938,69 @@ def test_import_star():
                                                             "from blah import _potato\n")
     assert SortImports(file_contents=test_input, combine_star=True).output == ("from blah import *\n")
 
+def test_include_trailing_comma():
+    """Test for the include_trailing_comma option"""
+    test_output_grid = SortImports(
+        file_contents=SHORT_IMPORT,
+        multi_line_output=WrapModes.GRID,
+        line_length=40,
+        include_trailing_comma=True,
+    ).output
+    assert test_output_grid == (
+        "from third_party import (lib1, lib2,\n"
+        "                         lib3, lib4,)\n"
+    )
+
+    test_output_vertical = SortImports(
+        file_contents=SHORT_IMPORT,
+        multi_line_output=WrapModes.VERTICAL,
+        line_length=40,
+        include_trailing_comma=True,
+    ).output
+    assert test_output_vertical == (
+        "from third_party import (lib1,\n"
+        "                         lib2,\n"
+        "                         lib3,\n"
+        "                         lib4,)\n"
+    )
+
+    test_output_vertical_indent = SortImports(
+        file_contents=SHORT_IMPORT,
+        multi_line_output=WrapModes.VERTICAL_HANGING_INDENT,
+        line_length=40,
+        include_trailing_comma=True,
+    ).output
+    assert test_output_vertical_indent == (
+        "from third_party import (\n"
+        "    lib1,\n"
+        "    lib2,\n"
+        "    lib3,\n"
+        "    lib4,\n"
+        ")\n"
+    )
+
+    test_output_vertical_grid = SortImports(
+        file_contents=SHORT_IMPORT,
+        multi_line_output=WrapModes.VERTICAL_GRID,
+        line_length=40,
+        include_trailing_comma=True,
+    ).output
+    assert test_output_vertical_grid == (
+        "from third_party import (\n"
+        "    lib1, lib2, lib3, lib4,)\n"
+    )
+
+    test_output_vertical_grid_grouped = SortImports(
+        file_contents=SHORT_IMPORT,
+        multi_line_output=WrapModes.VERTICAL_GRID_GROUPED,
+        line_length=40,
+        include_trailing_comma=True,
+    ).output
+    assert test_output_vertical_grid_grouped == (
+        "from third_party import (\n"
+        "    lib1, lib2, lib3, lib4,\n"
+        ")\n"
+    )
 
 def test_similar_to_std_library():
     """Test to ensure modules that are named similarly to a standard library import don't end up clobbered"""