force grid wrapping regardles of line length

3 files changed, 48 insertions, 27 deletions

diff --git a/isort/isort.py b/isort/isort.py
index b875cfc9..13dc6d44 100644
--- a/isort/isort.py
+++ b/isort/isort.py
@@ -381,33 +381,34 @@ class SortImports(object):
                     if not from_imports:
                         import_statement = ""
                     if len(import_statement) > self.config['line_length']:
-                        if len(from_imports) > 1:
-                            output_mode = settings.WrapModes._fields[self.config.get('multi_line_output',
-                                                                                        0)].lower()
-                            formatter = getattr(self, "_output_" + output_mode, self._output_grid)
-                            dynamic_indent = " " * (len(import_start) + 1)
-                            indent = self.config['indent']
-                            line_length = self.config['wrap_length'] or self.config['line_length']
-                            import_statement = formatter(import_start, copy.copy(from_imports),
-                                                        dynamic_indent, indent, line_length, comments)
-                            if self.config['balanced_wrapping']:
-                                lines = import_statement.split("\n")
-                                line_count = len(lines)
-                                if len(lines) > 1:
-                                    minimum_length = min([len(line) for line in lines[:-1]])
-                                else:
-                                    minimum_length = 0
-                                new_import_statement = import_statement
-                                while (len(lines[-1]) < minimum_length and
-                                        len(lines) == line_count and line_length > 10):
-                                    import_statement = new_import_statement
-                                    line_length -= 1
-                                    new_import_statement = formatter(import_start, copy.copy(from_imports),
-                                                                    dynamic_indent, indent, line_length, comments)
-                                    lines = new_import_statement.split("\n")
-                        else:
-                            import_statement = self._wrap(import_statement)
-
+                        import_statement = self._wrap(import_statement)
+                    if len(from_imports) > 1 and (
+                        len(import_statement) > self.config['line_length']
+                        or self.config.get('force_from_wrap')
+                    ):
+                        output_mode = settings.WrapModes._fields[self.config.get('multi_line_output',
+                                                                                    0)].lower()
+                        formatter = getattr(self, "_output_" + output_mode, self._output_grid)
+                        dynamic_indent = " " * (len(import_start) + 1)
+                        indent = self.config['indent']
+                        line_length = self.config['wrap_length'] or self.config['line_length']
+                        import_statement = formatter(import_start, copy.copy(from_imports),
+                                                    dynamic_indent, indent, line_length, comments)
+                        if self.config['balanced_wrapping']:
+                            lines = import_statement.split("\n")
+                            line_count = len(lines)
+                            if len(lines) > 1:
+                                minimum_length = min([len(line) for line in lines[:-1]])
+                            else:
+                                minimum_length = 0
+                            new_import_statement = import_statement
+                            while (len(lines[-1]) < minimum_length and
+                                    len(lines) == line_count and line_length > 10):
+                                import_statement = new_import_statement
+                                line_length -= 1
+                                new_import_statement = formatter(import_start, copy.copy(from_imports),
+                                                                dynamic_indent, indent, line_length, comments)
+                                lines = new_import_statement.split("\n")
 
                 if import_statement:
                     above_comments = self.comments['above']['from'].get(module, None)
diff --git a/isort/main.py b/isort/main.py
index 54d4f677..0f9b89d6 100755
--- a/isort/main.py
+++ b/isort/main.py
@@ -112,6 +112,8 @@ def main():
                         help="Switches the typical ordering preference, showing from imports first then straight ones.")
     parser.add_argument('-wl', '--wrap-length', dest='wrap_length',
                         help="Specifies how long lines that are wrapped should be, if not set line_length is used.")
+    parser.add_argument('-fgw', '--force-grid-wrap',  action='store_true', dest="force_grid_wrap",
+                        help='Force from imports to be grid wrapped regardless of line length')
 
     arguments = dict((key, value) for (key, value) in itemsview(vars(parser.parse_args())) if value)
     file_names = arguments.pop('files', [])
diff --git a/test_isort.py b/test_isort.py
index 1e906bee..457b9ef2 100644
--- a/test_isort.py
+++ b/test_isort.py
@@ -1260,3 +1260,21 @@ def test_consistency():
     """Ensures consistency of handling even when dealing with non ordered-by-type imports"""
     test_input = "from sqlalchemy.dialects.postgresql import ARRAY, array\n"
     assert SortImports(file_contents=test_input, order_by_type=True).output == test_input
+
+
+def test_force_grid_wrap():
+    """Ensures removing imports works as expected."""
+    test_input = ("from foo import lib6, lib7\n"
+                  "from bar import lib2\n")
+    test_output = SortImports(
+      file_contents=test_input,
+      force_from_wrap=True,
+      multi_line_output=WrapModes.VERTICAL_HANGING_INDENT
+      ).output
+    print(test_output)
+    assert test_output == """from bar import lib2
+from foo import (
+    lib6,
+    lib7
+)
+"""