features to align assignment operators and dictionary colons

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/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 ::=
@@ -309,6 +310,10 @@ def Visit_typedargslist(self, node):  # pylint: disable=invalid-name
         tname = True
         _SetArgListSubtype(child, subtypes.TYPED_NAME,
                            subtypes.TYPED_NAME_ARG_LIST)
+        # NOTE Every element of the tynamme argument
+        # should have this list type
+        _AppendSubtypeRec(child, subtypes.TYPED_NAME_ARG_LIST)
+
       elif child.type == grammar_token.COMMA:
         tname = False
       elif child.type == grammar_token.EQUAL and tname:

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,82 @@ 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_dict_colon(self):
+    # if the token is dictionary colon and
+    # the dictionary has no comp_for
+    return self.value == ':' and self.previous_token.is_dict_key
+
+  @property
+  def is_dict_key(self):
+    # if the token is dictionary key which is not preceded by doubel stars and
+    # the dictionary has no comp_for
+    return subtypes.DICTIONARY_KEY_PART in self.subtypes
+
+  @property
+  def is_dict_key_start(self):
+    # if the token is dictionary key start
+    return subtypes.DICTIONARY_KEY in self.subtypes
+
+  @property
+  def is_dict_value(self):
+    return subtypes.DICTIONARY_VALUE in self.subtypes
+
+  @property
+  def is_augassign(self):
+    augassigns = {'+=', '-=' , '*=' , '@=' , '/=' , '%=' , '&=' , '|=' , '^=' ,
+            '<<=' , '>>=' , '**=' , '//='}
+    return self.value in augassigns
+
+  @property
+  def is_argassign(self):
+     return (subtypes.DEFAULT_OR_NAMED_ASSIGN in self.subtypes
+            or subtypes.VARARGS_LIST in self.subtypes)
+
+  @property
+  def is_argname(self):
+    # it's the argument part before argument assignment operator,
+    # including tnames and data type
+    # not the assign operator,
+    # not the value after the assign operator
+
+    # argument without assignment is also included
+    # the token is arg part before '=' but not after '='
+    if self.is_argname_start:
+        return True
+
+    # exclude comment inside argument list
+    if not self.is_comment:
+      # the token is any element in typed arglist
+      if subtypes.TYPED_NAME_ARG_LIST in self.subtypes:
+        return True
+
+    return False
+
+  @property
+  def is_argname_start(self):
+    # return true if it's the start of every argument entry
+    previous_subtypes = {0}
+    if self.previous_token:
+      previous_subtypes = self.previous_token.subtypes
+
+    return (
+        (not self.is_comment
+        and subtypes.DEFAULT_OR_NAMED_ASSIGN not in self.subtypes
+        and subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST in self.subtypes
+        and subtypes.DEFAULT_OR_NAMED_ASSIGN not in previous_subtypes
+        and (not subtypes.PARAMETER_STOP in self.subtypes
+        or subtypes.PARAMETER_START in self.subtypes)
+        )
+        or # if there is comment, the arg after it is the argname start
+        (not self.is_comment and self.previous_token and self.previous_token.is_comment
+        and
+        (subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST in previous_subtypes
+        or subtypes.TYPED_NAME_ARG_LIST in self.subtypes
+        or subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST in self.subtypes))
+        )