fix handling of comments in from imports (#829)

Fixes #671

diff --git a/black.py b/black.py
index dd6e3720eebffb810ead12c79a7d9c7730cadc1e..9ecfbe18fe8443a6bb9c7640de036523127cad35 100644 (file)
--- a/black.py
+++ b/black.py
@@ -2405,10 +2405,17 @@ def bracket_split_build_line(
         if leaves:
             # Since body is a new indent level, remove spurious leading whitespace.
             normalize_prefix(leaves[0], inside_brackets=True)
-            # Ensure a trailing comma when expected.
+            # Ensure a trailing comma for imports, but be careful not to add one after
+            # any comments.
             if original.is_import:
-                if leaves[-1].type != token.COMMA:
-                    leaves.append(Leaf(token.COMMA, ","))
+                for i in range(len(leaves) - 1, -1, -1):
+                    if leaves[i].type == STANDALONE_COMMENT:
+                        continue
+                    elif leaves[i].type == token.COMMA:
+                        break
+                    else:
+                        leaves.insert(i + 1, Leaf(token.COMMA, ","))
+                        break
     # Populate the line
     for leaf in leaves:
         result.append(leaf, preformatted=True)

diff --git a/tests/data/comments7.py b/tests/data/comments7.py
new file mode 100644 (file)
index 0000000..4159d84
--- /dev/null
+++ b/tests/data/comments7.py
@@ -0,0 +1,51 @@
+from .config import (
+    Any,
+    Bool,
+    ConfigType,
+    ConfigTypeAttributes,
+    Int,
+    Path,
+    #  String,
+    #  resolve_to_config_type,
+    #  DEFAULT_TYPE_ATTRIBUTES,
+)
+
+
+from .config import (
+    Any,
+    Bool,
+    ConfigType,
+    ConfigTypeAttributes,
+    Int,
+    no_comma_here_yet
+    #  and some comments,
+    #  resolve_to_config_type,
+    #  DEFAULT_TYPE_ATTRIBUTES,
+)
+
+# output
+
+from .config import (
+    Any,
+    Bool,
+    ConfigType,
+    ConfigTypeAttributes,
+    Int,
+    Path,
+    #  String,
+    #  resolve_to_config_type,
+    #  DEFAULT_TYPE_ATTRIBUTES,
+)
+
+
+from .config import (
+    Any,
+    Bool,
+    ConfigType,
+    ConfigTypeAttributes,
+    Int,
+    no_comma_here_yet,
+    #  and some comments,
+    #  resolve_to_config_type,
+    #  DEFAULT_TYPE_ATTRIBUTES,
+)

diff --git a/tests/test_black.py b/tests/test_black.py
index 3e0b8a189bfcaa051f4056cc6246a6b56dd73a15..86175aad3041c331680ead5fc92dcabf72a0717b 100644 (file)
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -388,6 +388,14 @@ class BlackTestCase(unittest.TestCase):
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, black.FileMode())
 
+    @patch("black.dump_to_file", dump_to_stderr)
+    def test_comments7(self) -> None:
+        source, expected = read_data("comments7")
+        actual = fs(source)
+        self.assertFormatEqual(expected, actual)
+        black.assert_equivalent(source, actual)
+        black.assert_stable(source, actual, black.FileMode())
+
     @patch("black.dump_to_file", dump_to_stderr)
     def test_cantfit(self) -> None:
         source, expected = read_data("cantfit")