Repository: google/yapf
Pull request #820 (2020-03-04 01:07:23)

Knob for spaces around subscript colon
**SPACES_AROUND_SUBSCRIPT_COLON**

Use spaces around the subscript / slice operator.  For example:

```python
  my_list[1 : 10 : 2]
```

This Bool option (False by default) inserts spaces around colons when
they're used as a subscript / slice operator.

There's a potential interaction with `SPACE_INSIDE_BRACKETS`, so there's
an extra test case in TestsForSpacesAroundSubscriptColon to verify the
behaviour with both enabled.


Patch:
diff --git a/CHANGELOG b/CHANGELOG
index e1feffb..e24d0ee 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -13,6 +13,8 @@
 - Don't split between two-word comparison operators: "is not", "not in", etc.
 - New knob `SPACE_INSIDE_BRACKETS` to add spaces inside brackets, braces, and
   parentheses.
+- New knob `SPACES_AROUND_SUBSCRIPT_COLON` to add spaces around the subscript /
+  slice operator.
 
 ## [0.29.0] 2019-11-28
 ### Added
diff --git a/README.rst b/README.rst
index 1fd04ea..125a3b5 100644
--- a/README.rst
+++ b/README.rst
@@ -545,6 +545,13 @@ Knobs
     Set to ``True`` to prefer spaces around the assignment operator for default
     or keyword arguments.
 
+``SPACES_AROUND_SUBSCRIPT_COLON``
+    Use spaces around the subscript / slice operator.  For example:
+
+    .. code-block:: python
+
+        my_list[1 : 10 : 2]
+
 ``SPACES_BEFORE_COMMENT``
     The number of spaces required before a trailing comment.
     This can be a single value (representing the number of spaces
diff --git a/yapf/yapflib/format_token.py b/yapf/yapflib/format_token.py
index 79b2b5c..e6af942 100644
--- a/yapf/yapflib/format_token.py
+++ b/yapf/yapflib/format_token.py
@@ -313,6 +313,12 @@ class FormatToken(object):
     """Token is an operator in a simple expression."""
     return Subtype.SIMPLE_EXPRESSION in self.subtypes
 
+  @property
+  @py3compat.lru_cache()
+  def is_subscript_colon(self):
+    """Token is a subscript colon."""
+    return Subtype.SUBSCRIPT_COLON in self.subtypes
+
   @property
   @py3compat.lru_cache()
   def name(self):
diff --git a/yapf/yapflib/style.py b/yapf/yapflib/style.py
index 5f2483f..84994b3 100644
--- a/yapf/yapflib/style.py
+++ b/yapf/yapflib/style.py
@@ -219,6 +219,11 @@ _STYLE_HELP = dict(
       Use spaces around the power operator."""),
     SPACES_AROUND_DEFAULT_OR_NAMED_ASSIGN=textwrap.dedent("""\
       Use spaces around default or named assigns."""),
+    SPACES_AROUND_SUBSCRIPT_COLON=textwrap.dedent("""\
+      Use spaces around the subscript / slice operator.  For example:
+
+        my_list[1 : 10 : 2]
+      """),
     SPACES_BEFORE_COMMENT=textwrap.dedent("""\
       The number of spaces required before a trailing comment.
       This can be a single value (representing the number of spaces
@@ -394,6 +399,7 @@ def CreatePEP8Style():
       SPACE_INSIDE_BRACKETS=False,
       SPACES_AROUND_POWER_OPERATOR=False,
       SPACES_AROUND_DEFAULT_OR_NAMED_ASSIGN=False,
+      SPACES_AROUND_SUBSCRIPT_COLON=False,
       SPACES_BEFORE_COMMENT=2,
       SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED=False,
       SPLIT_ALL_COMMA_SEPARATED_VALUES=False,
@@ -577,6 +583,7 @@ _STYLE_OPTION_VALUE_CONVERTER = dict(
     SPACE_INSIDE_BRACKETS=_BoolConverter,
     SPACES_AROUND_POWER_OPERATOR=_BoolConverter,
     SPACES_AROUND_DEFAULT_OR_NAMED_ASSIGN=_BoolConverter,
+    SPACES_AROUND_SUBSCRIPT_COLON=_BoolConverter,
     SPACES_BEFORE_COMMENT=_IntOrIntListConverter,
     SPLIT_ARGUMENTS_WHEN_COMMA_TERMINATED=_BoolConverter,
     SPLIT_ALL_COMMA_SEPARATED_VALUES=_BoolConverter,
diff --git a/yapf/yapflib/unwrapped_line.py b/yapf/yapflib/unwrapped_line.py
index e4cb428..ec28f9b 100644
--- a/yapf/yapflib/unwrapped_line.py
+++ b/yapf/yapflib/unwrapped_line.py
@@ -261,6 +261,10 @@ def _PriorityIndicatingNoSpace(tok):
   return _HasPrecedence(tok)
 
 
+def _IsSubscriptColonAndValuePair(token1, token2):
+  return (token1.is_number or token1.is_name) and token2.is_subscript_colon
+
+
 def _SpaceRequiredBetween(left, right):
   """Return True if a space is required between the left and right token."""
   lval = left.value
@@ -293,6 +297,11 @@ def _SpaceRequiredBetween(left, right):
       return True
     if rval in pytree_utils.CLOSING_BRACKETS and lval == ':':
       return True
+  if (style.Get('SPACES_AROUND_SUBSCRIPT_COLON') and
+      (_IsSubscriptColonAndValuePair(left, right) or
+       _IsSubscriptColonAndValuePair(right, left))):
+    # Supersede the "never want a space before a colon or comma" check.
+    return True
   if rval in ':,':
     # Otherwise, we never want a space before a colon or comma.
     return False


Test patch:
diff --git a/yapftests/reformatter_pep8_test.py b/yapftests/reformatter_pep8_test.py
index 8f336e1..bdd074a 100644
--- a/yapftests/reformatter_pep8_test.py
+++ b/yapftests/reformatter_pep8_test.py
@@ -843,5 +843,77 @@ class TestsForSpacesInsideBrackets(yapf_test_helper.YAPFTest):
     self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
 
 
+class TestsForSpacesAroundSubscriptColon(yapf_test_helper.YAPFTest):
+  """Test the SPACES_AROUND_SUBSCRIPT_COLON style option."""
+  unformatted_code = textwrap.dedent("""\
+    a = list1[ : ]
+    b = list2[ slice_start: ]
+    c = list3[ slice_start:slice_end ]
+    d = list4[ slice_start:slice_end: ]
+    e = list5[ slice_start:slice_end:slice_step ]
+    a1 = list1[ : ]
+    b1 = list2[ 1: ]
+    c1 = list3[ 1:20 ]
+    d1 = list4[ 1:20: ]
+    e1 = list5[ 1:20:3 ]
+  """)
+
+  def testEnabled(self):
+    style.SetGlobalStyle(
+        style.CreateStyleFromConfig('{spaces_around_subscript_colon: True}'))
+    expected_formatted_code = textwrap.dedent("""\
+      a = list1[:]
+      b = list2[slice_start :]
+      c = list3[slice_start : slice_end]
+      d = list4[slice_start : slice_end :]
+      e = list5[slice_start : slice_end : slice_step]
+      a1 = list1[:]
+      b1 = list2[1 :]
+      c1 = list3[1 : 20]
+      d1 = list4[1 : 20 :]
+      e1 = list5[1 : 20 : 3]
+    """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(self.unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+  def testWithSpaceInsideBrackets(self):
+    style.SetGlobalStyle(
+        style.CreateStyleFromConfig('{'
+                                    'spaces_around_subscript_colon: true, '
+                                    'space_inside_brackets: true,'
+                                    '}'))
+    expected_formatted_code = textwrap.dedent("""\
+      a = list1[ : ]
+      b = list2[ slice_start : ]
+      c = list3[ slice_start : slice_end ]
+      d = list4[ slice_start : slice_end : ]
+      e = list5[ slice_start : slice_end : slice_step ]
+      a1 = list1[ : ]
+      b1 = list2[ 1 : ]
+      c1 = list3[ 1 : 20 ]
+      d1 = list4[ 1 : 20 : ]
+      e1 = list5[ 1 : 20 : 3 ]
+    """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(self.unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+  def testDefault(self):
+    style.SetGlobalStyle(style.CreatePEP8Style())
+    expected_formatted_code = textwrap.dedent("""\
+      a = list1[:]
+      b = list2[slice_start:]
+      c = list3[slice_start:slice_end]
+      d = list4[slice_start:slice_end:]
+      e = list5[slice_start:slice_end:slice_step]
+      a1 = list1[:]
+      b1 = list2[1:]
+      c1 = list3[1:20]
+      d1 = list4[1:20:]
+      e1 = list5[1:20:3]
+    """)
+    uwlines = yapf_test_helper.ParseAndUnwrap(self.unformatted_code)
+    self.assertCodeEqual(expected_formatted_code, reformatter.Reformat(uwlines))
+
+
 if __name__ == '__main__':
   unittest.main()