Feature align assignment

CHANGELOG

@@ -2,6 +2,10 @@
 # All notable changes to this project will be documented in this file.
 # This project adheres to [Semantic Versioning](http://semver.org/).
 
+## [0.41.1] 2022-08-30
+### Added
+- Add 4 new knobs to align assignment operators and dictionary colons. They are align_assignment, align_argument_assignment, align_dict_colon and new_alignment_after_commentline.
+
 ## [0.40.0] UNRELEASED
 ### Added
 - Add a new Python parser to generate logical lines.

CONTRIBUTORS

@@ -15,3 +15,4 @@ Sam Clegg <[email protected]>
 Łukasz Langa <[email protected]>
 Oleg Butuzov <[email protected]>
 Mauricio Herrera Cuadra <[email protected]>
+Xiao Wang <[email protected]>

README.rst

@@ -390,6 +390,61 @@ Options::
 Knobs
 =====
 
+``ALIGN_ASSIGNMENT``
+    Align assignment or augmented assignment operators.
+    If there is a blank line or a newline comment or a multiline object
+    (e.g. a dictionary, a list, a function call) in between,
+    it will start new block alignment. Lines in the same block have the same
+    indentation level.
+
+    .. code-block:: python
+
+        a   = 1
+        abc = 2
+        if condition == None:
+            var       += ''
+            var_long  -= 4
+        b  = 3
+        bc = 4
+
+``ALIGN_ARGUMENT_ASSIGNMENT``
+    Align assignment operators in the argument list if they are all split on newlines.
+    Arguments without assignment in between will initiate new block alignment calulation;
+    for example, a comment line.
+    Multiline objects in between will also initiate a new alignment block.
+
+    .. code-block:: python
+
+        rglist = test(
+            var_first  = 0,
+            var_second = '',
+            var_dict   = {
+                "key_1" : '',
+                "key_2" : 2,
+                "key_3" : True,
+            },
+            var_third     = 1,
+            var_very_long = None )
+
+``ALIGN_DICT_COLON``
+    Align the colons in the dictionary if all entries in dictionay are split on newlines
+    or 'EACH_DICT_ENTRY_ON_SEPERATE_LINE' is set True.
+    A commentline or multi-line object in between will start new alignment block.
+
+    .. code-block:: python
+
+        fields =
+            {
+                "field" : "ediid",
+                "type"  : "text",
+                # key: value
+                "required" : True,
+            }
+
+``NEW_ALIGNMENT_AFTER_COMMENTLINE``
+    Make it optional to start a new alignmetn block for assignment
+    alignment and colon alignment after a comment line.
+
 ``ALIGN_CLOSING_BRACKET_WITH_VISUAL_INDENT``
     Align closing bracket with visual indentation.
 

yapf/pytree/split_penalty.py

@@ -89,7 +89,7 @@ def Visit_classdef(self, node):  # pylint: disable=invalid-name
     if len(node.children) > 4:
       # opening '('
       _SetUnbreakable(node.children[2])
-    # ':'
+      # ':'
     _SetUnbreakable(node.children[-2])
     self.DefaultNodeVisit(node)
 

yapf/pytree/subtype_assigner.py

@@ -240,6 +240,7 @@ def Visit_argument(self, node):  # pylint: disable=invalid-name
     # argument ::=
     #     test [comp_for] | test '=' test
     self._ProcessArgLists(node)
+    #TODO add a subtype to each argument?
 
   def Visit_arglist(self, node):  # pylint: disable=invalid-name
     # arglist ::=

yapf/yapflib/format_decision_state.py

@@ -978,6 +978,7 @@ def _GetNewlineColumn(self):
         not self.param_list_stack[-1].SplitBeforeClosingBracket(
             top_of_stack.indent) and top_of_stack.indent
         == ((self.line.depth + 1) * style.Get('INDENT_WIDTH'))):
+      # NOTE: comment inside argument list is not excluded in subtype assigner
       if (subtypes.PARAMETER_START in current.subtypes or
           (previous.is_comment and
            subtypes.PARAMETER_START in previous.subtypes)):

yapf/yapflib/format_token.py

@@ -322,3 +322,15 @@ def is_pytype_comment(self):
   def is_copybara_comment(self):
     return self.is_comment and re.match(
         r'#.*\bcopybara:\s*(strip|insert|replace)', self.value)
+
+  @property
+  def is_assign(self):
+    return subtypes.ASSIGN_OPERATOR in self.subtypes
+
+  @property
+  def is_augassign(self):
+    augassigns = {
+        '+=', '-=', '*=', '@=', '/=', '%=', '&=', '|=', '^=', '<<=', '>>=',
+        '**=', '//='
+    }
+    return self.value in augassigns

yapf/yapflib/reformatter.py

@@ -22,6 +22,7 @@
 from __future__ import unicode_literals
 
 import collections
+from distutils.errors import LinkError
 import heapq
 import re
 
@@ -102,6 +103,9 @@ def Reformat(llines, verify=False, lines=None):
     final_lines.append(lline)
     prev_line = lline
 
+  if style.Get('ALIGN_ASSIGNMENT'):
+    _AlignAssignment(final_lines)
+
   _AlignTrailingComments(final_lines)
   return _FormatFinalLines(final_lines, verify)
 
@@ -394,6 +398,172 @@ def _AlignTrailingComments(final_lines):
       final_lines_index += 1
 
 
+def _AlignAssignment(final_lines):
+  """Align assignment operators and augmented assignment operators to the same column"""
+
+  final_lines_index = 0
+  while final_lines_index < len(final_lines):
+    line = final_lines[final_lines_index]
+
+    assert line.tokens
+    process_content = False
+
+    for tok in line.tokens:
+      if tok.is_assign or tok.is_augassign:
+        # all pre assignment variable lengths in one block of lines
+        all_pa_variables_lengths = []
+        max_variables_length = 0
+
+        while True:
+          # EOF
+          if final_lines_index + len(all_pa_variables_lengths) == len(
+              final_lines):
+            break
+
+          this_line_index = final_lines_index + len(all_pa_variables_lengths)
+          this_line = final_lines[this_line_index]
+
+          next_line = None
+          if this_line_index < len(final_lines) - 1:
+            next_line = final_lines[final_lines_index +
+                                    len(all_pa_variables_lengths) + 1]
+
+          assert this_line.tokens, next_line.tokens
+
+          # align them differently when there is a blank line in between
+          if (all_pa_variables_lengths and
+              this_line.tokens[0].formatted_whitespace_prefix.startswith('\n\n')
+             ):
+            break
+
+          # if there is a standalone comment or keyword statement line
+          # or other lines without assignment in between, break
+          elif (all_pa_variables_lengths and True not in [
+              tok.is_assign or tok.is_augassign for tok in this_line.tokens
+          ]):
+            if this_line.tokens[0].is_comment:
+              if style.Get('NEW_ALIGNMENT_AFTER_COMMENTLINE'):
+                break
+            else:
+              break
+
+          if this_line.disable:
+            all_pa_variables_lengths.append([])
+            continue
+
+          variables_content = ''
+          pa_variables_lengths = []
+          contain_object = False
+          line_tokens = this_line.tokens
+          # only one assignment expression is on each line
+          for index in range(len(line_tokens)):
+            line_tok = line_tokens[index]
+
+            prefix = line_tok.formatted_whitespace_prefix
+            newline_index = prefix.rfind('\n')
+            if newline_index != -1:
+              variables_content = ''
+              prefix = prefix[newline_index + 1:]
+
+            if line_tok.is_assign or line_tok.is_augassign:
+              next_toks = [
+                  line_tokens[i] for i in range(index + 1, len(line_tokens))
+              ]
+              # if there is object(list/tuple/dict) with newline entries, break,
+              # update the alignment so far and start to calulate new alignment
+              for tok in next_toks:
+                if tok.value in ['(', '[', '{'] and tok.next_token:
+                  if (tok.next_token.formatted_whitespace_prefix.startswith(
+                      '\n') or
+                      (tok.next_token.is_comment and tok.next_token.next_token
+                       .formatted_whitespace_prefix.startswith('\n'))):
+                    pa_variables_lengths.append(len(variables_content))
+                    contain_object = True
+                    break
+              if not contain_object:
+                if line_tok.is_assign:
+                  pa_variables_lengths.append(len(variables_content))
+                # if augassign, add the extra augmented part to the max length caculation
+                elif line_tok.is_augassign:
+                  pa_variables_lengths.append(
+                      len(variables_content) + len(line_tok.value) - 1)
+              # don't add the tokens
+              # after the assignment operator
+              break
+            else:
+              variables_content += '{}{}'.format(prefix, line_tok.value)
+
+          if pa_variables_lengths:
+            max_variables_length = max(max_variables_length,
+                                       max(pa_variables_lengths))
+
+          all_pa_variables_lengths.append(pa_variables_lengths)
+
+          # after saving this line's max variable length,
+          # we check if next line has the same depth as this line,
+          # if not, we don't want to calculate their max variable length together
+          # so we break the while loop, update alignment so far, and
+          # then go to next line that has '='
+          if next_line:
+            if this_line.depth != next_line.depth:
+              break
+          # if this line contains objects with newline entries,
+          # start new block alignment
+          if contain_object:
+            break
+
+        # if no update of max_length, just go to the next block
+        if max_variables_length == 0:
+          continue
+
+        max_variables_length += 2
+
+        # Update the assignment token values based on the max variable length
+        for all_pa_variables_lengths_index, pa_variables_lengths in enumerate(
+            all_pa_variables_lengths):
+          if not pa_variables_lengths:
+            continue
+          this_line = final_lines[final_lines_index +
+                                  all_pa_variables_lengths_index]
+
+          # only the first assignment operator on each line
+          pa_variables_lengths_index = 0
+          for line_tok in this_line.tokens:
+            if line_tok.is_assign or line_tok.is_augassign:
+              assert pa_variables_lengths[0] < max_variables_length
+
+              if pa_variables_lengths_index < len(pa_variables_lengths):
+                whitespace = ' ' * (
+                    max_variables_length - pa_variables_lengths[0] - 1)
+
+                assign_content = '{}{}'.format(whitespace,
+                                               line_tok.value.strip())
+
+                existing_whitespace_prefix = \
+                    line_tok.formatted_whitespace_prefix.lstrip('\n')
+
+                # in case the existing spaces are larger than padded spaces
+                if (len(whitespace) == 1 or len(whitespace) > 1 and
+                    len(existing_whitespace_prefix) > len(whitespace)):
+                  line_tok.whitespace_prefix = ''
+                elif assign_content.startswith(existing_whitespace_prefix):
+                  assign_content = assign_content[len(existing_whitespace_prefix
+                                                     ):]
+
+                # update the assignment operator value
+                line_tok.value = assign_content
+
+              pa_variables_lengths_index += 1
+
+        final_lines_index += len(all_pa_variables_lengths)
+
+        process_content = True
+        break
+
+    if not process_content:
+      final_lines_index += 1
+
+
 def _FormatFinalLines(final_lines, verify):
   """Compose the final output from the finalized lines."""
   formatted_code = []

yapf/yapflib/style.py

@@ -54,6 +54,13 @@ def SetGlobalStyle(style):
 _STYLE_HELP = dict(
     ALIGN_CLOSING_BRACKET_WITH_VISUAL_INDENT=textwrap.dedent("""\
       Align closing bracket with visual indentation."""),
+    ALIGN_ASSIGNMENT=textwrap.dedent("""\
+      Align assignment or augmented assignment operators.
+      If there is a blank line or newline comment or objects with newline entries in between,
+      it will start new block alignment."""),
+    NEW_ALIGNMENT_AFTER_COMMENTLINE=textwrap.dedent("""\
+      Start new assignment or colon alignment when there is a newline comment in between."""
+                                                   ),
     ALLOW_MULTILINE_LAMBDAS=textwrap.dedent("""\
       Allow lambdas to be formatted on more than one line."""),
     ALLOW_MULTILINE_DICTIONARY_KEYS=textwrap.dedent("""\
@@ -419,6 +426,8 @@ def CreatePEP8Style():
   """Create the PEP8 formatting style."""
   return dict(
       ALIGN_CLOSING_BRACKET_WITH_VISUAL_INDENT=True,
+      ALIGN_ASSIGNMENT=False,
+      NEW_ALIGNMENT_AFTER_COMMENTLINE=False,
       ALLOW_MULTILINE_LAMBDAS=False,
       ALLOW_MULTILINE_DICTIONARY_KEYS=False,
       ALLOW_SPLIT_BEFORE_DEFAULT_OR_NAMED_ASSIGNS=True,
@@ -607,6 +616,8 @@ def _IntOrIntListConverter(s):
 # Note: this dict has to map all the supported style options.
 _STYLE_OPTION_VALUE_CONVERTER = dict(
     ALIGN_CLOSING_BRACKET_WITH_VISUAL_INDENT=_BoolConverter,
+    ALIGN_ASSIGNMENT=_BoolConverter,
+    NEW_ALIGNMENT_AFTER_COMMENTLINE=_BoolConverter,
     ALLOW_MULTILINE_LAMBDAS=_BoolConverter,
     ALLOW_MULTILINE_DICTIONARY_KEYS=_BoolConverter,
     ALLOW_SPLIT_BEFORE_DEFAULT_OR_NAMED_ASSIGNS=_BoolConverter,

yapftests/format_token_test.py

@@ -15,10 +15,11 @@
 
 import unittest
 
-from lib2to3 import pytree
+from lib2to3 import pytree, pygram
 from lib2to3.pgen2 import token
 
 from yapf.yapflib import format_token
+from yapf.pytree import subtype_assigner
 
 
 class TabbedContinuationAlignPaddingTest(unittest.TestCase):