Wrap arithmetic and binary arithmetic expressions in invisible parentheses (#2001)

[etc/vim.git] / src / black / __init__.py

diff --git a/src/black/__init__.py b/src/black/__init__.py
index 07251d55fb543169ff34ef29e3993e73d0473f26..6919468609ce34c6079a27c0794120b54c57acf2 100644 (file)
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -2049,6 +2049,8 @@ class LineGenerator(Visitor[Line]):
 
     def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
         """Visit a statement without nested statements."""
+        if first_child_is_arith(node):
+            wrap_in_parentheses(node, node.children[0], visible=False)
         is_suite_like = node.parent and node.parent.type in STATEMENT
         if is_suite_like:
             if self.mode.is_pyi and is_stub_body(node):
@@ -5613,6 +5615,17 @@ def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
     return wrapped
 
 
+def first_child_is_arith(node: Node) -> bool:
+    """Whether first child is an arithmetic or a binary arithmetic expression"""
+    expr_types = {
+        syms.arith_expr,
+        syms.shift_expr,
+        syms.xor_expr,
+        syms.and_expr,
+    }
+    return bool(node.children and node.children[0].type in expr_types)
+
+
 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
     """Wrap `child` in parentheses.