Feature align assignment (continuation of PR#1030)

CHANGELOG.md

@@ -4,6 +4,46 @@
 
 ## (0.41.0) UNRELEASED
 ### Added
+- New `ALIGN_ASSIGNMENT` and `ALIGN_ASSIGNMENT_RESTART_AFTER_COMMENTS` flags.
+ `ALIGN_ASSIGNMENT` is a new knob that enables alignment of assignment 
+ operators in consecutive lines.
+
+ A blank line or multiline object (function, multiline dictionary, ...) will 
+ always interrupt the alignment block and start a new one. Comments can also
+ optionally interrupt the block, a behaviour that can be turned on with
+ `ALIGN_ASSIGNMENT_RESTART_AFTER_COMMENTS`.
+
+ For example, with `ALIGN_ASSIGNMENT` off:
+  ```      
+  val_first      = 1
+  val_second += 2
+  # comment
+  val_third    = 3
+  ```
+
+  will be formatted as:
+  ```      
+  val_first = 1
+  val_second += 2
+  # comment
+  val_third = 3
+  ```
+  whereas with `ALIGN_ASSIGNMENT` on and 
+  `ALIGN_ASSIGNMENT_RESTART_AFTER_COMMENTS` off we get:
+  ```
+  val_first   = 1
+  val_second += 2
+  # comment
+  val_third   = 3
+  ```
+  while if they are both on:
+  ```    
+  val_first   = 1
+  val_second += 2
+  # comment
+  val_third = 3
+  ```
+
 - New `DISABLE_SPLIT_LIST_WITH_COMMENT` flag.
  `DISABLE_SPLIT_LIST_WITH_COMMENT` is a new knob that changes the
   behavior of splitting a list when a comment is present inside the list.

CONTRIBUTORS

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

README.md

@@ -357,6 +357,64 @@ optional arguments:
 
 ## Knobs
 
+#### `ALIGN_ASSIGNMENT`
+
+>    Align assignment or augmented assignment operators. If there is a blank
+>    line or objects with newline entries in between, it will start new block
+>    alignment. For exemple:
+
+```python
+    val_first = 1
+    val_second += 2
+    object = {
+        entry1:1,
+        entry2:2,
+        entry3:3,
+        }
+    val_third = 3
+```
+
+>    will be formatted as follows:
+
+```python
+    val_first   = 1
+    val_second += 2
+    object      = {
+        entry1: 1,
+        entry2: 2,
+        entry3: 3,
+    }
+    val_third = 3
+```
+
+#### `ALIGN_ASSIGNMENT_RESTART_AFTER_COMMENTS`
+
+>    Allows to choose whether to align assignments before and after a comment
+>    line independently. For exemple, if off:
+
+```python  
+  val_first      = 1
+  val_second += 2
+  # comment
+  val_third    = 3
+```
+
+will be aligned as:
+```python      
+  val_first   = 1
+  val_second += 2
+  # comment
+  val_third   = 3
+```
+
+whereas if we turn it on:
+```python
+  val_first   = 1
+  val_second += 2
+  # comment
+  val_third = 3
+```
+
 #### `ALIGN_CLOSING_BRACKET_WITH_VISUAL_INDENT`
 
 >    Align closing bracket with visual indentation.

yapf/yapflib/format_token.py

@@ -327,3 +327,26 @@ 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
+  @lru_cache()
+  def is_augassign(self):
+    return self.value in frozenset({
+        '+=',
+        '-=',
+        '*=',
+        '@=',
+        '/=',
+        '//=',
+        '%=',
+        '<<=',
+        '>>=',
+        '|=',
+        '&=',
+        '^=',
+        '**=',
+    })