Use a set for subtypes.

This prevents us from adding the same subtype multiple times.

2 files changed, 57 insertions, 54 deletions

diff --git a/yapf/yapflib/format_token.py b/yapf/yapflib/format_token.py
index af31d32..42eeaa5 100644
--- a/yapf/yapflib/format_token.py
+++ b/yapf/yapflib/format_token.py
@@ -120,7 +120,7 @@ class FormatToken(object):
 
     stypes = pytree_utils.GetNodeAnnotation(node,
                                             pytree_utils.Annotation.SUBTYPE)
-    self.subtypes = [subtypes.NONE] if stypes is None else stypes
+    self.subtypes = {subtypes.NONE} if not stypes else stypes
     self.is_pseudo = hasattr(node, 'is_pseudo') and node.is_pseudo
 
   @property
@@ -215,6 +215,9 @@ class FormatToken(object):
   def ClosesScope(self):
     return self.value in pytree_utils.CLOSING_BRACKETS
 
+  def AddSubtype(self, subtype):
+    self.subtypes.add(subtype)
+
   def __repr__(self):
     msg = ('FormatToken(name={0}, value={1}, column={2}, lineno={3}, '
            'splitpenalty={4}'.format(
diff --git a/yapftests/subtype_assigner_test.py b/yapftests/subtype_assigner_test.py
index 2438646..66be68e 100644
--- a/yapftests/subtype_assigner_test.py
+++ b/yapftests/subtype_assigner_test.py
@@ -52,7 +52,7 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('def', [subtypes.NONE]),
+            ('def', {subtypes.NONE}),
             ('foo', {subtypes.FUNC_DEF}),
             ('(', {subtypes.NONE}),
             ('a', {
@@ -92,15 +92,15 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
                 subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST,
             }),
             (')', {subtypes.NONE}),
-            (':', [subtypes.NONE]),
+            (':', {subtypes.NONE}),
         ],
         [
-            ('return', [subtypes.NONE]),
+            ('return', {subtypes.NONE}),
             ('-', {subtypes.UNARY_OPERATOR}),
-            ('x', [subtypes.NONE]),
+            ('x', {subtypes.NONE}),
             ('[', {subtypes.SUBSCRIPT_BRACKET}),
             (':', {subtypes.SUBSCRIPT_COLON}),
-            ('42', [subtypes.NONE]),
+            ('42', {subtypes.NONE}),
             (']', {subtypes.SUBSCRIPT_BRACKET}),
         ],
     ])
@@ -112,8 +112,8 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('foo', [subtypes.NONE]),
-            ('(', [subtypes.NONE]),
+            ('foo', {subtypes.NONE}),
+            ('(', {subtypes.NONE}),
             ('x', {
                 subtypes.NONE,
                 subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST,
@@ -125,7 +125,7 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
             }),
             ('=', {subtypes.DEFAULT_OR_NAMED_ASSIGN}),
             ("'hello world'", {subtypes.NONE}),
-            (')', [subtypes.NONE]),
+            (')', {subtypes.NONE}),
         ],
     ])
 
@@ -137,7 +137,7 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('def', [subtypes.NONE]),
+            ('def', {subtypes.NONE}),
             ('foo', {subtypes.FUNC_DEF}),
             ('(', {subtypes.NONE}),
             ('strs', {
@@ -146,11 +146,11 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
                 subtypes.PARAMETER_STOP,
             }),
             (')', {subtypes.NONE}),
-            (':', [subtypes.NONE]),
+            (':', {subtypes.NONE}),
         ],
         [
-            ('return', [subtypes.NONE]),
-            ('{', [subtypes.NONE]),
+            ('return', {subtypes.NONE}),
+            ('{', {subtypes.NONE}),
             ('s', {subtypes.COMP_EXPR}),
             ('.', {subtypes.COMP_EXPR}),
             ('lower', {subtypes.COMP_EXPR}),
@@ -163,7 +163,7 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
             ('s', {subtypes.COMP_FOR}),
             ('in', {subtypes.COMP_FOR}),
             ('strs', {subtypes.COMP_FOR}),
-            ('}', [subtypes.NONE]),
+            ('}', {subtypes.NONE}),
         ],
     ])
 
@@ -174,7 +174,7 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines,
                                    [[('not', {subtypes.UNARY_OPERATOR}),
-                                     ('a', [subtypes.NONE])]])
+                                     ('a', {subtypes.NONE})]])
 
   def testBitwiseOperators(self):
     code = textwrap.dedent("""\
@@ -183,25 +183,25 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('x', [subtypes.NONE]),
+            ('x', {subtypes.NONE}),
             ('=', {subtypes.ASSIGN_OPERATOR}),
-            ('(', [subtypes.NONE]),
-            ('(', [subtypes.NONE]),
-            ('a', [subtypes.NONE]),
+            ('(', {subtypes.NONE}),
+            ('(', {subtypes.NONE}),
+            ('a', {subtypes.NONE}),
             ('|', {subtypes.BINARY_OPERATOR}),
-            ('(', [subtypes.NONE]),
-            ('b', [subtypes.NONE]),
+            ('(', {subtypes.NONE}),
+            ('b', {subtypes.NONE}),
             ('^', {subtypes.BINARY_OPERATOR}),
-            ('3', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('3', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('&', {subtypes.BINARY_OPERATOR}),
-            ('c', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('c', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('<<', {subtypes.BINARY_OPERATOR}),
-            ('3', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('3', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('>>', {subtypes.BINARY_OPERATOR}),
-            ('1', [subtypes.NONE]),
+            ('1', {subtypes.NONE}),
         ],
     ])
 
@@ -212,37 +212,37 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('x', [subtypes.NONE]),
+            ('x', {subtypes.NONE}),
             ('=', {subtypes.ASSIGN_OPERATOR}),
-            ('(', [subtypes.NONE]),
-            ('(', [subtypes.NONE]),
-            ('a', [subtypes.NONE]),
+            ('(', {subtypes.NONE}),
+            ('(', {subtypes.NONE}),
+            ('a', {subtypes.NONE}),
             ('+', {subtypes.BINARY_OPERATOR}),
-            ('(', [subtypes.NONE]),
-            ('b', [subtypes.NONE]),
+            ('(', {subtypes.NONE}),
+            ('b', {subtypes.NONE}),
             ('-', {
                 subtypes.BINARY_OPERATOR,
                 subtypes.SIMPLE_EXPRESSION,
             }),
-            ('3', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('3', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('*', {subtypes.BINARY_OPERATOR}),
-            ('(', [subtypes.NONE]),
-            ('1', [subtypes.NONE]),
+            ('(', {subtypes.NONE}),
+            ('1', {subtypes.NONE}),
             ('%', {
                 subtypes.BINARY_OPERATOR,
                 subtypes.SIMPLE_EXPRESSION,
             }),
-            ('c', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('c', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('@', {subtypes.BINARY_OPERATOR}),
-            ('d', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('d', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('/', {subtypes.BINARY_OPERATOR}),
-            ('3', [subtypes.NONE]),
-            (')', [subtypes.NONE]),
+            ('3', {subtypes.NONE}),
+            (')', {subtypes.NONE}),
             ('//', {subtypes.BINARY_OPERATOR}),
-            ('1', [subtypes.NONE]),
+            ('1', {subtypes.NONE}),
         ],
     ])
 
@@ -253,13 +253,13 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('x', [subtypes.NONE]),
+            ('x', {subtypes.NONE}),
             ('[', {subtypes.SUBSCRIPT_BRACKET}),
-            ('0', [subtypes.NONE]),
+            ('0', {subtypes.NONE}),
             (':', {subtypes.SUBSCRIPT_COLON}),
-            ('42', [subtypes.NONE]),
+            ('42', {subtypes.NONE}),
             (':', {subtypes.SUBSCRIPT_COLON}),
-            ('1', [subtypes.NONE]),
+            ('1', {subtypes.NONE}),
             (']', {subtypes.SUBSCRIPT_BRACKET}),
         ],
     ])
@@ -271,15 +271,15 @@ class SubtypeAssignerTest(yapf_test_helper.YAPFTest):
     uwlines = yapf_test_helper.ParseAndUnwrap(code)
     self._CheckFormatTokenSubtypes(uwlines, [
         [
-            ('[', [subtypes.NONE]),
-            ('a', [subtypes.NONE]),
-            (',', [subtypes.NONE]),
+            ('[', {subtypes.NONE}),
+            ('a', {subtypes.NONE}),
+            (',', {subtypes.NONE}),
             ('*', {
                 subtypes.UNARY_OPERATOR,
                 subtypes.VARARGS_STAR,
             }),
-            ('b', [subtypes.NONE]),
-            (']', [subtypes.NONE]),
+            ('b', {subtypes.NONE}),
+            (']', {subtypes.NONE}),
         ],
     ])