Repository: google/yapf
Pull request #649 (2018-11-19 01:26:01)

Updated trailing comment horizontal alignment based on a list of potential values
This addresses https://github.com/google/yapf/issues/541.

Thanks!

gwelymernans: Could you squash these commits? It's hard to review this way.

Patch:
diff --git a/README.rst b/README.rst
index f5773e0..493b4ed 100644
--- a/README.rst
+++ b/README.rst
@@ -484,6 +484,41 @@ Knobs
 
 ``SPACES_BEFORE_COMMENT``
     The number of spaces required before a trailing comment.
+    This can be a single value (representing the number of spaces
+    before each trailing comment) or list of of values (representing
+    alignment column values; trailing comments within a block will
+    be aligned to the first column value that is greater than the maximum
+    line length within the block). For example:
+    
+    With spaces_before_comment=5:
+    
+      1 + 1 # Adding values
+    
+    will be formatted as:
+    
+      1 + 1     # Adding values <-- 5 spaces between the end of the statement and comment
+    
+    With spaces_before_comment=15, 20:
+    
+      1 + 1 # Adding values
+      two + two # More adding
+    
+      longer_statement # This is a longer statement
+      short # This is a shorter statement
+    
+      a_very_long_statement_that_extends_beyond_the_final_column # Comment
+      short # This is a shorter statement
+    
+    will be formatted as:
+    
+      1 + 1          # Adding values <-- end of line comments in block aligned to col 15
+      two + two      # More adding
+    
+      longer_statement    # This is a longer statement <-- end of line comments in block aligned to col 20
+      short               # This is a shorter statement
+    
+      a_very_long_statement_that_extends_beyond_the_final_column  # Comment <-- the end of line comments are aligned based on the line length
+      short                                                       # This is a shorter statement
 
 ``SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET``
     Insert a space between the ending comma and closing bracket of a list, etc.
diff --git a/yapf/yapflib/format_decision_state.py b/yapf/yapflib/format_decision_state.py
index b08235e..27ff540 100644
--- a/yapf/yapflib/format_decision_state.py
+++ b/yapf/yapflib/format_decision_state.py
@@ -522,6 +522,11 @@ class FormatDecisionState(object):
     previous = current.previous_token
 
     spaces = current.spaces_required_before
+    if isinstance(spaces, list):
+      # Don't set the value here, as we need to look at the lines near
+      # this one to determine the actual horizontal alignment value.
+      spaces = 0
+
     if not dry_run:
       current.AddWhitespacePrefix(newlines_before=0, spaces=spaces)
 
@@ -742,7 +747,11 @@ class FormatDecisionState(object):
     previous = current.previous_token
     top_of_stack = self.stack[-1]
 
-    if current.spaces_required_before > 2 or self.line.disable:
+    if isinstance(current.spaces_required_before, list):
+      # Don't set the value here, as we need to look at the lines near
+      # this one to determine the actual horizontal alignment value.
+      return 0
+    elif current.spaces_required_before > 2 or self.line.disable:
       return current.spaces_required_before
 
     if current.OpensScope():
diff --git a/yapf/yapflib/reformatter.py b/yapf/yapflib/reformatter.py
index 7d26038..23c40dd 100644
--- a/yapf/yapflib/reformatter.py
+++ b/yapf/yapflib/reformatter.py
@@ -95,6 +95,8 @@ def Reformat(uwlines, verify=False, lines=None):
 
     final_lines.append(uwline)
     prev_uwline = uwline
+
+  _AlignTrailingComments(final_lines)
   return _FormatFinalLines(final_lines, verify)
 
 
@@ -246,6 +248,134 @@ def _CanPlaceOnSingleLine(uwline):
           not any(tok.is_comment for tok in uwline.tokens[:last_index]))
 
 
+def _AlignTrailingComments(final_lines):
+  final_lines_index = 0
+  while final_lines_index < len(final_lines):
+    line = final_lines[final_lines_index]
+    assert line.tokens
+
+    processed_content = False
+
+    for tok in line.tokens:
+      if (tok.is_comment and isinstance(tok.spaces_required_before, list) and
+          tok.value.startswith('#')):
+        # All trailing comments and comments that appear on a line by themselves
+        # in this block should be indented at the same level. The block is
+        # terminated by an empty line or EOF. Enumerate through each line in
+        # the block and calculate the max line length. Once complete, use the
+        # first col value greater than that value and create the necessary for
+        # each line accordingly.
+        all_pc_line_lengths = []  # All pre-comment line lengths
+        max_line_length = 0
+
+        while True:
+          # EOF
+          if final_lines_index + len(all_pc_line_lengths) == len(final_lines):
+            break
+
+          this_line = final_lines[final_lines_index + len(all_pc_line_lengths)]
+
+          # Blank line - note that content is preformatted so we don't need to
+          # worry about spaces/tabs; a blank line will always be '\n\n'.
+          assert this_line.tokens
+          if (all_pc_line_lengths and
+              this_line.tokens[0].formatted_whitespace_prefix.startswith('\n\n')
+             ):
+            break
+
+          if this_line.disable:
+            all_pc_line_lengths.append([])
+            continue
+
+          # Calculate the length of each line in this unwrapped line.
+          line_content = ''
+          pc_line_lengths = []
+
+          for line_tok in this_line.tokens:
+            whitespace_prefix = line_tok.formatted_whitespace_prefix
+
+            newline_index = whitespace_prefix.rfind('\n')
+            if newline_index != -1:
+              max_line_length = max(max_line_length, len(line_content))
+              line_content = ''
+
+              whitespace_prefix = whitespace_prefix[newline_index + 1:]
+
+            if line_tok.is_comment:
+              pc_line_lengths.append(len(line_content))
+            else:
+              line_content += "{}{}".format(whitespace_prefix, line_tok.value)
+
+          if pc_line_lengths:
+            max_line_length = max(max_line_length, max(pc_line_lengths))
+
+          all_pc_line_lengths.append(pc_line_lengths)
+
+        # Calculate the aligned column value
+        max_line_length += 2
+
+        aligned_col = None
+        for potential_col in tok.spaces_required_before:
+          if potential_col > max_line_length:
+            aligned_col = potential_col
+            break
+
+        if aligned_col is None:
+          aligned_col = max_line_length
+
+        # Update the comment token values based on the aligned values
+        for all_pc_line_lengths_index, pc_line_lengths in enumerate(
+            all_pc_line_lengths):
+          if not pc_line_lengths:
+            continue
+
+          this_line = final_lines[final_lines_index + all_pc_line_lengths_index]
+
+          pc_line_length_index = 0
+          for line_tok in this_line.tokens:
+            if line_tok.is_comment:
+              assert pc_line_length_index < len(pc_line_lengths)
+              assert pc_line_lengths[pc_line_length_index] < aligned_col
+
+              # Note that there may be newlines embedded in the comments, so
+              # we need to apply a whitespace prefix to each line.
+              whitespace = ' ' * (
+                  aligned_col - pc_line_lengths[pc_line_length_index] - 1)
+              pc_line_length_index += 1
+
+              line_content = []
+
+              for comment_line_index, comment_line in enumerate(
+                  line_tok.value.split('\n')):
+                line_content.append("{}{}".format(whitespace,
+                                                  comment_line.strip()))
+
+                if comment_line_index == 0:
+                  whitespace = ' ' * (aligned_col - 1)
+
+              line_content = '\n'.join(line_content)
+
+              # Account for initial whitespace already slated for the
+              # beginning of the line.
+              existing_whitespace_prefix = \
+                line_tok.formatted_whitespace_prefix.lstrip('\n')
+
+              if line_content.startswith(existing_whitespace_prefix):
+                line_content = line_content[len(existing_whitespace_prefix):]
+
+              line_tok.value = line_content
+
+          assert pc_line_length_index == len(pc_line_lengths)
+
+        final_lines_index += len(all_pc_line_lengths)
+
+        processed_content = True
+        break
+
+    if not processed_content:
+      final_lines_index += 1
+
+
 def _FormatFinalLines(final_lines, verify):
   """Compose the final output from the finalized lines."""
   formatted_code = []
diff --git a/yapf/yapflib/style.py b/yapf/yapflib/style.py
index 87f1b54..0fe5982 100644
--- a/yapf/yapflib/style.py
+++ b/yapf/yapflib/style.py
@@ -176,7 +176,44 @@ _STYLE_HELP = dict(
     SPACES_AROUND_DEFAULT_OR_NAMED_ASSIGN=textwrap.dedent("""\
       Use spaces around default or named assigns."""),
     SPACES_BEFORE_COMMENT=textwrap.dedent("""\
-      The number of spaces required before a trailing comment."""),
+      The number of spaces required before a trailing comment.
+      This can be a single value (representing the number of spaces
+      before each trailing comment) or list of of values (representing
+      alignment column values; trailing comments within a block will
+      be aligned to the first column value that is greater than the maximum
+      line length within the block). For example:
+      
+      With spaces_before_comment=5:
+
+        1 + 1 # Adding values
+
+      will be formatted as:
+
+        1 + 1     # Adding values <-- 5 spaces between the end of the statement and comment
+
+      With spaces_before_comment=15, 20:
+
+        1 + 1 # Adding values
+        two + two # More adding
+
+        longer_statement # This is a longer statement
+        short # This is a shorter statement
+
+        a_very_long_statement_that_extends_beyond_the_final_column # Comment
+        short # This is a shorter statement
+
+      will be formatted as:
+
+        1 + 1          # Adding values <-- end of line comments in block aligned to col 15
+        two + two      # More adding
+
+        longer_statement    # This is a longer statement <-- end of line comments in block aligned to col 20
+        short               # This is a shorter statement
+
+        a_very_long_statement_that_extends_beyond_the_final_column  # Comment <-- the end of line comments are aligned based on the line length
+        short                                                       # This is a shorter statement
+
+      """),
     SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED=textwrap.dedent("""\
       Split before arguments if the argument list is terminated by a
       comma."""),
@@ -420,6 +457,20 @@ def _BoolConverter(s):
   return py3compat.CONFIGPARSER_BOOLEAN_STATES[s.lower()]
 
 
+def _IntListConverter(s):
+  """Option value converter for a comma-separated list of integers."""
+  s = s.strip()
+  if s.startswith('[') and s.endswith(']'):
+    s = s[1:-1]
+
+  return [int(part.strip()) for part in s.split(',') if part.strip()]
+
+
+def _IntOrIntListConverter(s):
+  """Option value converter for an integer or list of integers."""
+  return _IntListConverter(s) if ',' in s else int(s)
+
+
 # Different style options need to have their values interpreted differently when
 # read from the config file. This dict maps an option name to a "converter"
 # function that accepts the string read for the option's value from the file and
@@ -453,7 +504,7 @@ _STYLE_OPTION_VALUE_CONVERTER = dict(
     SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET=_BoolConverter,
     SPACES_AROUND_POWER_OPERATOR=_BoolConverter,
     SPACES_AROUND_DEFAULT_OR_NAMED_ASSIGN=_BoolConverter,
-    SPACES_BEFORE_COMMENT=int,
+    SPACES_BEFORE_COMMENT=_IntOrIntListConverter,
     SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED=_BoolConverter,
     SPLIT_ALL_COMMA_SEPARATED_VALUES=_BoolConverter,
     SPLIT_BEFORE_BITWISE_OPERATOR=_BoolConverter,
diff --git a/yapf/yapflib/unwrapped_line.py b/yapf/yapflib/unwrapped_line.py
index 0bf4fba..f37f3e4 100644
--- a/yapf/yapflib/unwrapped_line.py
+++ b/yapf/yapflib/unwrapped_line.py
@@ -73,7 +73,19 @@ class UnwrappedLine(object):
         token.spaces_required_before = 1
 
       tok_len = len(token.value) if not token.is_pseudo_paren else 0
-      token.total_length = prev_length + tok_len + token.spaces_required_before
+
+      spaces_required_before = token.spaces_required_before
+      if isinstance(spaces_required_before, list):
+        assert token.is_comment, token
+
+        # If here, we are looking at a comment token that appears on a line
+        # with other tokens (but because it is a comment, it is always the last token).
+        # Rather than specifying the actual number of spaces here, hard code a value of
+        # 0 and then set it later. This logic only works because this comment token is
+        # guaranteed to be the last token in the list.
+        spaces_required_before = 0
+
+      token.total_length = prev_length + tok_len + spaces_required_before
 
       # The split penalty has to be computed before {must|can}_break_before,
       # because these may use it for their decision.


Test patch:
diff --git a/yapftests/style_test.py b/yapftests/style_test.py
index 7f4a465..f525a48 100644
--- a/yapftests/style_test.py
+++ b/yapftests/style_test.py
@@ -53,6 +53,15 @@ class UtilsTest(unittest.TestCase):
     self.assertEqual(style._BoolConverter('false'), False)
     self.assertEqual(style._BoolConverter('0'), False)
 
+  def testIntListConverter(self):
+    self.assertEqual(style._IntListConverter("1, 2, 3"), [1, 2, 3])
+    self.assertEqual(style._IntListConverter("[ 1, 2, 3 ]"), [1, 2, 3])
+    self.assertEqual(style._IntListConverter("[ 1, 2, 3, ]"), [1, 2, 3])
+
+  def testIntOrIntListConverter(self):
+    self.assertEqual(style._IntOrIntListConverter("10"), 10)
+    self.assertEqual(style._IntOrIntListConverter("1, 2, 3"), [1, 2, 3])
+
 
 def _LooksLikeChromiumStyle(cfg):
   return (cfg['INDENT_WIDTH'] == 2 and
diff --git a/yapftests/yapf_test.py b/yapftests/yapf_test.py
index c9302aa..24a1adc 100644
--- a/yapftests/yapf_test.py
+++ b/yapftests/yapf_test.py
@@ -1436,5 +1436,289 @@ class DiffIndentTest(unittest.TestCase):
     self._Check(unformatted_code, expected_formatted_code)
 
 
+class HorizontallyAlignedTrailingCommentsTest(unittest.TestCase):
+
+  @staticmethod
+  def _OwnStyle():
+    my_style = style.CreatePEP8Style()
+    my_style["SPACES_BEFORE_COMMENT"] = [
+        15,
+        25,
+        35,
+    ]
+    return my_style
+
+  def _Check(self, unformatted_code, expected_formatted_code):
+    formatted_code, _ = yapf_api.FormatCode(
+        unformatted_code, style_config=style.SetGlobalStyle(self._OwnStyle()))
+    self.assertEqual(expected_formatted_code, formatted_code)
+
+  def setUp(self):
+    self.maxDiff = None
+
+  def testSimple(self):
+    unformatted_code = textwrap.dedent("""\
+        foo = '1' # Aligned at first list value
+
+        foo = '2__<15>' # Aligned at second list value
+
+        foo = '3____________<25>' # Aligned at third list value
+
+        foo = '4______________________<35>' # Aligned beyond list values
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        foo = '1'     # Aligned at first list value
+        
+        foo = '2__<15>'         # Aligned at second list value
+        
+        foo = '3____________<25>'         # Aligned at third list value
+        
+        foo = '4______________________<35>' # Aligned beyond list values
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlock(self):
+    unformatted_code = textwrap.dedent("""\
+        func(1)     # Line 1
+        func(2) # Line 2
+        # Line 3
+        func(3)                             # Line 4
+                                            # Line 5
+                                            # Line 6
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        func(1)       # Line 1
+        func(2)       # Line 2
+                      # Line 3
+        func(3)       # Line 4
+                      # Line 5
+                      # Line 6
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockWithLongLine(self):
+    unformatted_code = textwrap.dedent("""\
+        func(1)     # Line 1
+        func___________________(2) # Line 2
+        # Line 3
+        func(3)                             # Line 4
+                                            # Line 5
+                                            # Line 6
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        func(1)                           # Line 1
+        func___________________(2)        # Line 2
+                                          # Line 3
+        func(3)                           # Line 4
+                                          # Line 5
+                                          # Line 6
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockFuncSuffix(self):
+    unformatted_code = textwrap.dedent("""\
+        func(1)     # Line 1
+        func(2) # Line 2
+        # Line 3
+        func(3)                             # Line 4
+                                        # Line 5 - SpliceComments makes this a new block
+                                    # Line 6
+
+        def Func():
+            pass
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        func(1)       # Line 1
+        func(2)       # Line 2
+                      # Line 3
+        func(3)       # Line 4
+
+        # Line 5 - SpliceComments makes this a new block
+        # Line 6
+
+
+        def Func():
+            pass
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockCommentSuffix(self):
+    unformatted_code = textwrap.dedent("""\
+        func(1)     # Line 1
+        func(2) # Line 2
+        # Line 3
+        func(3)                             # Line 4
+                                        # Line 5 - SpliceComments makes this part of the previous block
+                                    # Line 6
+
+                                            # Aligned with prev comment block
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        func(1)       # Line 1
+        func(2)       # Line 2
+                      # Line 3
+        func(3)       # Line 4
+                      # Line 5 - SpliceComments makes this part of the previous block
+                      # Line 6
+
+                      # Aligned with prev comment block
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockIndentedFuncSuffix(self):
+    unformatted_code = textwrap.dedent("""\
+        if True:
+            func(1)     # Line 1
+            func(2) # Line 2
+            # Line 3
+            func(3)                             # Line 4
+                                                # Line 5 - SpliceComments makes this a new block
+                                                # Line 6
+
+                                                # Aligned with Func
+
+            def Func():
+                pass
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        if True:
+            func(1)   # Line 1
+            func(2)   # Line 2
+                      # Line 3
+            func(3)   # Line 4
+                      
+            # Line 5 - SpliceComments makes this a new block
+            # Line 6
+
+            # Aligned with Func
+
+
+            def Func():
+                pass
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockIndentedCommentSuffix(self):
+    unformatted_code = textwrap.dedent("""\
+        if True:
+            func(1)     # Line 1
+            func(2) # Line 2
+            # Line 3
+            func(3)                             # Line 4
+                                                # Line 5
+                                                # Line 6
+            
+                                                # Not aligned
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        if True:
+            func(1)   # Line 1
+            func(2)   # Line 2
+                      # Line 3
+            func(3)   # Line 4
+                      # Line 5
+                      # Line 6
+
+            # Not aligned
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testBlockMultiIndented(self):
+    unformatted_code = textwrap.dedent("""\
+        if True:
+            if True:
+                if True:
+                    func(1)     # Line 1
+                    func(2) # Line 2
+                    # Line 3
+                    func(3)                             # Line 4
+                                                        # Line 5
+                                                        # Line 6
+                    
+                                                        # Not aligned
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        if True:
+            if True:
+                if True:
+                    func(1)     # Line 1
+                    func(2)     # Line 2
+                                # Line 3
+                    func(3)     # Line 4
+                                # Line 5
+                                # Line 6
+        
+                    # Not aligned
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testArgs(self):
+    unformatted_code = textwrap.dedent("""\
+        def MyFunc(
+            arg1,   # Desc 1
+            arg2,   # Desc 2
+            a_longer_var_name,  # Desc 3
+            arg4,
+            arg5,   # Desc 5
+            arg6,
+        ):
+            pass
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        def MyFunc(
+                arg1,                     # Desc 1
+                arg2,                     # Desc 2
+                a_longer_var_name,        # Desc 3
+                arg4,
+                arg5,                     # Desc 5
+                arg6,
+        ):
+            pass
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testDisableBlock(self):
+    unformatted_code = textwrap.dedent("""\
+        a() # comment 1
+        b() # comment 2
+
+        # yapf: disable
+        c() # comment 3
+        d()   # comment 4
+        # yapf: enable
+
+        e() # comment 5
+        f() # comment 6
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        a()           # comment 1
+        b()           # comment 2
+
+        # yapf: disable
+        c() # comment 3
+        d()   # comment 4
+        # yapf: enable
+
+        e()           # comment 5
+        f()           # comment 6
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+  def testDisabledLine(self):
+    unformatted_code = textwrap.dedent("""\
+        short # comment 1
+        do_not_touch1 # yapf: disable
+        do_not_touch2   # yapf: disable
+        a_longer_statement # comment 2
+        """)
+    expected_formatted_code = textwrap.dedent("""\
+        short                   # comment 1
+        do_not_touch1 # yapf: disable
+        do_not_touch2   # yapf: disable
+        a_longer_statement      # comment 2
+        """)
+    self._Check(unformatted_code, expected_formatted_code)
+
+
 if __name__ == '__main__':
   unittest.main()