fix: bug where the doublestar operation had inconsistent formatting. (#4154)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

CHANGES.md

@@ -16,6 +16,8 @@
 
 - Move the `hug_parens_with_braces_and_square_brackets` feature to the unstable style
   due to an outstanding crash and proposed formatting tweaks (#4198)
+- Fixed a bug where base expressions caused inconsistent formatting of \*\* in tenary
+  expression (#4154)
 - Checking for newline before adding one on docstring that is almost at the line limit
   (#4185)
 

docs/the_black_code_style/future_style.md

@@ -28,6 +28,8 @@ Currently, the following features are included in the preview style:
   longer normalized
 - `typed_params_trailing_comma`: consistently add trailing commas to typed function
   parameters
+- `is_simple_lookup_for_doublestar_expression`: fix line length computation for certain
+  expressions that involve the power operator
 - `docstring_check_for_newline`: checks if there is a newline before the terminating
   quotes of a docstring
 

src/black/mode.py

@@ -177,6 +177,7 @@ class Preview(Enum):
     wrap_long_dict_values_in_parens = auto()
     multiline_string_handling = auto()
     typed_params_trailing_comma = auto()
+    is_simple_lookup_for_doublestar_expression = auto()
     docstring_check_for_newline = auto()
 
 

src/black/resources/black.schema.json

@@ -86,6 +86,7 @@
           "wrap_long_dict_values_in_parens",
           "multiline_string_handling",
           "typed_params_trailing_comma",
+          "is_simple_lookup_for_doublestar_expression",
           "docstring_check_for_newline"
         ]
       },

src/black/trans.py

@@ -29,7 +29,7 @@
 
 from black.comments import contains_pragma_comment
 from black.lines import Line, append_leaves
-from black.mode import Feature, Mode
+from black.mode import Feature, Mode, Preview
 from black.nodes import (
     CLOSING_BRACKETS,
     OPENING_BRACKETS,
@@ -94,43 +94,36 @@ def hug_power_op(
     else:
         raise CannotTransform("No doublestar token was found in the line.")
 
-    def is_simple_lookup(index: int, step: Literal[1, -1]) -> bool:
+    def is_simple_lookup(index: int, kind: Literal[1, -1]) -> bool:
         # Brackets and parentheses indicate calls, subscripts, etc. ...
         # basically stuff that doesn't count as "simple". Only a NAME lookup
         # or dotted lookup (eg. NAME.NAME) is OK.
-        if step == -1:
+        if Preview.is_simple_lookup_for_doublestar_expression not in mode:
-            disallowed = {token.RPAR, token.RSQB}
+            return original_is_simple_lookup_func(line, index, kind)
+
         else:
-            disallowed = {token.LPAR, token.LSQB}
+            if kind == -1:
+                return handle_is_simple_look_up_prev(
+                    line, index, {token.RPAR, token.RSQB}
+                )
+            else:
+                return handle_is_simple_lookup_forward(
+                    line, index, {token.LPAR, token.LSQB}
+                )
 
-        while 0 <= index < len(line.leaves):
+    def is_simple_operand(index: int, kind: Literal[1, -1]) -> bool:
-            current = line.leaves[index]
-            if current.type in disallowed:
-                return False
-            if current.type not in {token.NAME, token.DOT} or current.value == "for":
-                # If the current token isn't disallowed, we'll assume this is simple as
-                # only the disallowed tokens are semantically attached to this lookup
-                # expression we're checking. Also, stop early if we hit the 'for' bit
-                # of a comprehension.
-                return True
-
-            index += step
-
-        return True
-
-    def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
         # An operand is considered "simple" if's a NAME, a numeric CONSTANT, a simple
         # lookup (see above), with or without a preceding unary operator.
         start = line.leaves[index]
         if start.type in {token.NAME, token.NUMBER}:
-            return is_simple_lookup(index, step=(1 if kind == "exponent" else -1))
+            return is_simple_lookup(index, kind)
 
         if start.type in {token.PLUS, token.MINUS, token.TILDE}:
             if line.leaves[index + 1].type in {token.NAME, token.NUMBER}:
-                # step is always one as bases with a preceding unary op will be checked
+                # kind is always one as bases with a preceding unary op will be checked
                 # for simplicity starting from the next token (so it'll hit the check
                 # above).
-                return is_simple_lookup(index + 1, step=1)
+                return is_simple_lookup(index + 1, kind=1)
 
         return False
 
@@ -145,9 +138,9 @@ def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
         should_hug = (
             (0 < idx < len(line.leaves) - 1)
             and leaf.type == token.DOUBLESTAR
-            and is_simple_operand(idx - 1, kind="base")
+            and is_simple_operand(idx - 1, kind=-1)
             and line.leaves[idx - 1].value != "lambda"
-            and is_simple_operand(idx + 1, kind="exponent")
+            and is_simple_operand(idx + 1, kind=1)
         )
         if should_hug:
             new_leaf.prefix = ""
@@ -162,6 +155,99 @@ def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
     yield new_line
 
 
+def original_is_simple_lookup_func(
+    line: Line, index: int, step: Literal[1, -1]
+) -> bool:
+    if step == -1:
+        disallowed = {token.RPAR, token.RSQB}
+    else:
+        disallowed = {token.LPAR, token.LSQB}
+
+    while 0 <= index < len(line.leaves):
+        current = line.leaves[index]
+        if current.type in disallowed:
+            return False
+        if current.type not in {token.NAME, token.DOT} or current.value == "for":
+            # If the current token isn't disallowed, we'll assume this is
+            # simple as only the disallowed tokens are semantically
+            # attached to this lookup expression we're checking. Also,
+            # stop early if we hit the 'for' bit of a comprehension.
+            return True
+
+        index += step
+
+    return True
+
+
+def handle_is_simple_look_up_prev(line: Line, index: int, disallowed: Set[int]) -> bool:
+    """
+    Handling the determination of is_simple_lookup for the lines prior to the doublestar
+    token. This is required because of the need to isolate the chained expression
+    to determine the bracket or parenthesis belong to the single expression.
+    """
+    contains_disallowed = False
+    chain = []
+
+    while 0 <= index < len(line.leaves):
+        current = line.leaves[index]
+        chain.append(current)
+        if not contains_disallowed and current.type in disallowed:
+            contains_disallowed = True
+        if not is_expression_chained(chain):
+            return not contains_disallowed
+
+        index -= 1
+
+    return True
+
+
+def handle_is_simple_lookup_forward(
+    line: Line, index: int, disallowed: Set[int]
+) -> bool:
+    """
+    Handling decision is_simple_lookup for the lines behind the doublestar token.
+    This function is simplified to keep consistent with the prior logic and the forward
+    case are more straightforward and do not need to care about chained expressions.
+    """
+    while 0 <= index < len(line.leaves):
+        current = line.leaves[index]
+        if current.type in disallowed:
+            return False
+        if current.type not in {token.NAME, token.DOT} or (
+            current.type == token.NAME and current.value == "for"
+        ):
+            # If the current token isn't disallowed, we'll assume this is simple as
+            # only the disallowed tokens are semantically attached to this lookup
+            # expression we're checking. Also, stop early if we hit the 'for' bit
+            # of a comprehension.
+            return True
+
+        index += 1
+
+    return True
+
+
+def is_expression_chained(chained_leaves: List[Leaf]) -> bool:
+    """
+    Function to determine if the variable is a chained call.
+    (e.g., foo.lookup, foo().lookup, (foo.lookup())) will be recognized as chained call)
+    """
+    if len(chained_leaves) < 2:
+        return True
+
+    current_leaf = chained_leaves[-1]
+    past_leaf = chained_leaves[-2]
+
+    if past_leaf.type == token.NAME:
+        return current_leaf.type in {token.DOT}
+    elif past_leaf.type in {token.RPAR, token.RSQB}:
+        return current_leaf.type in {token.RSQB, token.RPAR}
+    elif past_leaf.type in {token.LPAR, token.LSQB}:
+        return current_leaf.type in {token.NAME, token.LPAR, token.LSQB}
+    else:
+        return False
+
+
 class StringTransformer(ABC):
     """
     An implementation of the Transformer protocol that relies on its

tests/data/cases/is_simple_lookup_for_doublestar_expression.py

@@ -0,0 +1,14 @@
+# flags: --preview
+m2 = None if not isinstance(dist, Normal) else m** 2 + s * 2
+m3 = None if not isinstance(dist, Normal) else m ** 2 + s * 2
+m4 = None if not isinstance(dist, Normal) else m**2 + s * 2
+m5 = obj.method(another_obj.method()).attribute **2
+m6 = None if ... else m**2 + s**2
+
+
+# output
+m2 = None if not isinstance(dist, Normal) else m**2 + s * 2
+m3 = None if not isinstance(dist, Normal) else m**2 + s * 2
+m4 = None if not isinstance(dist, Normal) else m**2 + s * 2
+m5 = obj.method(another_obj.method()).attribute ** 2
+m6 = None if ... else m**2 + s**2