From 9a6c88c7f4b4db14631f8dc9e8d44b3aed9d57c9 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=C5=81ukasz=20Langa?= <lukasz@langa.pl>
Date: Mon, 21 May 2018 16:35:45 -0700
Subject: [PATCH] Fix invalid code on stars in long from-imports being wrapped
 in parentheses

Fixes #234
---
 README.md               |  2 ++
 black.py                | 44 +++++++++++++++--------------------------
 tests/fmtonoff.py       |  8 ++++++++
 tests/import_spacing.py |  2 ++
 4 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/README.md b/README.md
index 9b42f20..90731d2 100644
--- a/README.md
+++ b/README.md
@@ -661,6 +661,8 @@ More details can be found in [CONTRIBUTING](CONTRIBUTING.md).
 
 * fixed optional parentheses being removed within `# fmt: off` sections (#224)
 
+* fixed invalid code produced when stars in very long imports were incorrectly 
+  wrapped in optional parentheses (#234)
 
 ### 18.5b0
 
diff --git a/black.py b/black.py
index 9fbacc1..afc37d9 100644
--- a/black.py
+++ b/black.py
@@ -1375,32 +1375,6 @@ class LineGenerator(Visitor[Line]):
             yield from self.line()
             yield from self.visit(child)
 
-    def visit_import_from(self, node: Node) -> Iterator[Line]:
-        """Visit import_from and maybe put invisible parentheses.
-
-        This is separate from `visit_stmt` because import statements don't
-        support arbitrary atoms and thus handling of parentheses is custom.
-        """
-        check_lpar = False
-        for index, child in enumerate(node.children):
-            if check_lpar:
-                if child.type == token.LPAR:
-                    # make parentheses invisible
-                    child.value = ""  # type: ignore
-                    node.children[-1].value = ""  # type: ignore
-                else:
-                    # insert invisible parentheses
-                    node.insert_child(index, Leaf(token.LPAR, ""))
-                    node.append_child(Leaf(token.RPAR, ""))
-                break
-
-            check_lpar = (
-                child.type == token.NAME and child.value == "import"  # type: ignore
-            )
-
-        for child in node.children:
-            yield from self.visit(child)
-
     def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
         """Remove a semicolon and put the other statement on a separate line."""
         yield from self.line()
@@ -1447,6 +1421,7 @@ class LineGenerator(Visitor[Line]):
         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
         self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
+        self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
         self.visit_async_funcdef = self.visit_async_stmt
         self.visit_decorated = self.visit_decorators
 
@@ -2343,7 +2318,7 @@ def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
         return  # This `node` has a prefix with `# fmt: off`, don't mess with parens.
 
     check_lpar = False
-    for child in list(node.children):
+    for index, child in enumerate(list(node.children)):
         if check_lpar:
             if child.type == syms.atom:
                 maybe_make_parens_invisible_in_atom(child)
@@ -2351,8 +2326,21 @@ def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
                 # wrap child in visible parentheses
                 lpar = Leaf(token.LPAR, "(")
                 rpar = Leaf(token.RPAR, ")")
-                index = child.remove() or 0
+                child.remove()
                 node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
+            elif node.type == syms.import_from:
+                # "import from" nodes store parentheses directly as part of
+                # the statement
+                if child.type == token.LPAR:
+                    # make parentheses invisible
+                    child.value = ""  # type: ignore
+                    node.children[-1].value = ""  # type: ignore
+                elif child.type != token.STAR:
+                    # insert invisible parentheses
+                    node.insert_child(index, Leaf(token.LPAR, ""))
+                    node.append_child(Leaf(token.RPAR, ""))
+                break
+
             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
                 # wrap child in invisible parentheses
                 lpar = Leaf(token.LPAR, "")
diff --git a/tests/fmtonoff.py b/tests/fmtonoff.py
index e76e6a7..5666dd8 100644
--- a/tests/fmtonoff.py
+++ b/tests/fmtonoff.py
@@ -6,6 +6,10 @@ from third_party import X, Y, Z
 
 from library import some_connection, \
                     some_decorator
+# fmt: off
+from third_party import (X,
+                         Y, Z)
+# fmt: on
 f'trigger 3.6 mode'
 # fmt: off
 def func_no_args():
@@ -102,6 +106,10 @@ from third_party import X, Y, Z
 
 from library import some_connection, some_decorator
 
+# fmt: off
+from third_party import (X,
+                         Y, Z)
+# fmt: on
 f"trigger 3.6 mode"
 # fmt: off
 def func_no_args():
diff --git a/tests/import_spacing.py b/tests/import_spacing.py
index cefa1e9..0c98650 100644
--- a/tests/import_spacing.py
+++ b/tests/import_spacing.py
@@ -24,6 +24,7 @@ from some_library import (
     Just, Enough, Libraries, To, Fit, In, This, Nice, Split, Which, We, No, Longer, Use
 )
 from name_of_a_company.extremely_long_project_name.component.ttypes import CuteLittleServiceHandlerFactoryyy
+from name_of_a_company.extremely_long_project_name.extremely_long_component_name.ttypes import *
 
 from .a.b.c.subprocess import *
 from . import (tasks)
@@ -87,6 +88,7 @@ from some_library import (
 from name_of_a_company.extremely_long_project_name.component.ttypes import (
     CuteLittleServiceHandlerFactoryyy
 )
+from name_of_a_company.extremely_long_project_name.extremely_long_component_name.ttypes import *
 
 from .a.b.c.subprocess import *
 from . import tasks
-- 
2.39.5