X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/7811f957f3f6e5a5daa67a3f36a867cbb6210c8d..c9b73b99ac567ce30c36596371bf0a1fad8191c9:/black.py

diff --git a/black.py b/black.py
index 36380e6..c2363e4 100644
--- a/black.py
+++ b/black.py
@@ -599,6 +599,22 @@ TEST_DESCENDANTS = {
     syms.term,
     syms.power,
 }
+ASSIGNMENTS = {
+    "=",
+    "+=",
+    "-=",
+    "*=",
+    "@=",
+    "/=",
+    "%=",
+    "&=",
+    "|=",
+    "^=",
+    "<<=",
+    ">>=",
+    "**=",
+    "//=",
+}
 COMPREHENSION_PRIORITY = 20
 COMMA_PRIORITY = 18
 TERNARY_PRIORITY = 16
@@ -631,8 +647,8 @@ class BracketTracker:
     bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = Factory(dict)
     delimiters: Dict[LeafID, Priority] = Factory(dict)
     previous: Optional[Leaf] = None
-    _for_loop_variable: bool = False
-    _lambda_arguments: bool = False
+    _for_loop_variable: int = 0
+    _lambda_arguments: int = 0
 
     def mark(self, leaf: Leaf) -> None:
         """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
@@ -694,7 +710,7 @@ class BracketTracker:
         """
         if leaf.type == token.NAME and leaf.value == "for":
             self.depth += 1
-            self._for_loop_variable = True
+            self._for_loop_variable += 1
             return True
 
         return False
@@ -703,7 +719,7 @@ class BracketTracker:
         """See `maybe_increment_for_loop_variable` above for explanation."""
         if self._for_loop_variable and leaf.type == token.NAME and leaf.value == "in":
             self.depth -= 1
-            self._for_loop_variable = False
+            self._for_loop_variable -= 1
             return True
 
         return False
@@ -716,7 +732,7 @@ class BracketTracker:
         """
         if leaf.type == token.NAME and leaf.value == "lambda":
             self.depth += 1
-            self._lambda_arguments = True
+            self._lambda_arguments += 1
             return True
 
         return False
@@ -725,7 +741,7 @@ class BracketTracker:
         """See `maybe_increment_lambda_arguments` above for explanation."""
         if self._lambda_arguments and leaf.type == token.COLON:
             self.depth -= 1
-            self._lambda_arguments = False
+            self._lambda_arguments -= 1
             return True
 
         return False
@@ -994,8 +1010,9 @@ class Line:
             and subscript_start.type == syms.subscriptlist
         ):
             subscript_start = child_towards(subscript_start, leaf)
-        return subscript_start is not None and any(
-            n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
+        return (
+            subscript_start is not None
+            and any(n.type in TEST_DESCENDANTS for n in subscript_start.pre_order())
         )
 
     def __str__(self) -> str:
@@ -1252,14 +1269,13 @@ class LineGenerator(Visitor[Line]):
         """Visit a statement.
 
         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
-        `def`, `with`, `class`, and `assert`.
+        `def`, `with`, `class`, `assert` and assignments.
 
         The relevant Python language `keywords` for a given statement will be
         NAME leaves within it. This methods puts those on a separate line.
 
-        `parens` holds pairs of nodes where invisible parentheses should be put.
-        Keys hold nodes after which opening parentheses should be put, values
-        hold nodes before which closing parentheses should be put.
+        `parens` holds a set of string leaf values immeditely after which
+        invisible parens should be put.
         """
         normalize_invisible_parens(node, parens_after=parens)
         for child in node.children:
@@ -1369,6 +1385,8 @@ class LineGenerator(Visitor[Line]):
         self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
         self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
         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_async_funcdef = self.visit_async_stmt
         self.visit_decorated = self.visit_decorators
 
@@ -2205,6 +2223,9 @@ def normalize_string_quotes(leaf: Leaf) -> None:
 def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
     """Make existing optional parentheses invisible or create new ones.
 
+    `parens_after` is a set of string leaf values immeditely after which parens
+    should be put.
+
     Standardizes on visible parentheses for single-element tuples, and keeps
     existing visible parentheses for other tuples and generator expressions.
     """