Compare the node's type instead of its name

5 files changed, 37 insertions, 32 deletions

diff --git a/yapf/yapflib/blank_line_calculator.py b/yapf/yapflib/blank_line_calculator.py
index 3741348..3d78646 100644
--- a/yapf/yapflib/blank_line_calculator.py
+++ b/yapf/yapflib/blank_line_calculator.py
@@ -22,6 +22,8 @@ Annotations:
   newlines: The number of newlines required before the node.
 """
 
+from lib2to3.pgen2 import token as grammar_token
+
 from yapf.yapflib import py3compat
 from yapf.yapflib import pytree_utils
 from yapf.yapflib import pytree_visitor
@@ -64,7 +66,7 @@ class _BlankLineCalculator(pytree_visitor.PyTreeVisitor):
 
   def Visit_simple_stmt(self, node):  # pylint: disable=invalid-name
     self.DefaultNodeVisit(node)
-    if pytree_utils.NodeName(node.children[0]) == 'COMMENT':
+    if node.children[0].type == grammar_token.COMMENT:
       self.last_comment_lineno = node.children[0].lineno
 
   def Visit_decorator(self, node):  # pylint: disable=invalid-name
@@ -174,4 +176,4 @@ def _StartsInZerothColumn(node):
 
 def _AsyncFunction(node):
   return (py3compat.PY3 and node.prev_sibling and
-          pytree_utils.NodeName(node.prev_sibling) == 'ASYNC')
+          node.prev_sibling.type == grammar_token.ASYNC)
diff --git a/yapf/yapflib/identify_container.py b/yapf/yapflib/identify_container.py
index 5c5fc5b..888dbbb 100644
--- a/yapf/yapflib/identify_container.py
+++ b/yapf/yapflib/identify_container.py
@@ -19,6 +19,8 @@ to the opening bracket and vice-versa.
   IdentifyContainers(): the main function exported by this module.
 """
 
+from lib2to3.pgen2 import token as grammar_token
+
 from yapf.yapflib import pytree_utils
 from yapf.yapflib import pytree_visitor
 
@@ -42,7 +44,7 @@ class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
 
     if len(node.children) != 3:
       return
-    if pytree_utils.NodeName(node.children[0]) != 'LPAR':
+    if node.children[0].type != grammar_token.LPAR:
       return
 
     if pytree_utils.NodeName(node.children[1]) == 'arglist':
@@ -59,7 +61,7 @@ class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
 
     if len(node.children) != 3:
       return
-    if pytree_utils.NodeName(node.children[0]) != 'LPAR':
+    if node.children[0].type != grammar_token.LPAR:
       return
 
     for child in node.children[1].children:
diff --git a/yapf/yapflib/pytree_unwrapper.py b/yapf/yapflib/pytree_unwrapper.py
index c7e32a3..88272ef 100644
--- a/yapf/yapflib/pytree_unwrapper.py
+++ b/yapf/yapflib/pytree_unwrapper.py
@@ -205,7 +205,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
     for child in node.children:
       index += 1
       self.Visit(child)
-      if pytree_utils.NodeName(child) == 'ASYNC':
+      if child.type == grammar_token.ASYNC:
         break
     for child in node.children[index].children:
       self.Visit(child)
@@ -221,18 +221,17 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
     for child in node.children:
       index += 1
       self.Visit(child)
-      if pytree_utils.NodeName(child) == 'ASYNC':
+      if child.type == grammar_token.ASYNC:
         break
     for child in node.children[index].children:
-      if pytree_utils.NodeName(child) == 'NAME' and child.value == 'else':
+      if child.type == grammar_token.NAME and child.value == 'else':
         self._StartNewLine()
       self.Visit(child)
 
   def Visit_decorator(self, node):  # pylint: disable=invalid-name
     for child in node.children:
       self.Visit(child)
-      if (pytree_utils.NodeName(child) == 'COMMENT' and
-          child == node.children[0]):
+      if child.type == grammar_token.COMMENT and child == node.children[0]:
         self._StartNewLine()
 
   def Visit_decorators(self, node):  # pylint: disable=invalid-name
@@ -386,8 +385,7 @@ def _DetermineMustSplitAnnotation(node):
   if not _ContainsComments(node):
     token = next(node.parent.leaves())
     if token.value == '(':
-      if sum(1 for ch in node.children
-             if pytree_utils.NodeName(ch) == 'COMMA') < 2:
+      if sum(1 for ch in node.children if ch.type == grammar_token.COMMA) < 2:
         return
     if (not isinstance(node.children[-1], pytree.Leaf) or
         node.children[-1].value != ','):
diff --git a/yapf/yapflib/split_penalty.py b/yapf/yapflib/split_penalty.py
index d4c3dc4..40aedb4 100644
--- a/yapf/yapflib/split_penalty.py
+++ b/yapf/yapflib/split_penalty.py
@@ -16,6 +16,7 @@
 import re
 
 from lib2to3 import pytree
+from lib2to3.pgen2 import token as grammar_token
 
 from yapf.yapflib import format_token
 from yapf.yapflib import py3compat
@@ -123,7 +124,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     allow_multiline_lambdas = style.Get('ALLOW_MULTILINE_LAMBDAS')
     if not allow_multiline_lambdas:
       for child in node.children:
-        if pytree_utils.NodeName(child) == 'COMMENT':
+        if child.type == grammar_token.COMMENT:
           if re.search(r'pylint:.*disable=.*\bg-long-lambda', child.value):
             allow_multiline_lambdas = True
             break
@@ -145,7 +146,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
 
   def Visit_arglist(self, node):  # pylint: disable=invalid-name
     # arglist ::= argument (',' argument)* [',']
-    if pytree_utils.NodeName(node.children[0]) == 'STAR':
+    if node.children[0].type == grammar_token.STAR:
       # Python 3 treats a star expression as a specific expression type.
       # Process it in that method.
       self.Visit_star_expr(node)
@@ -200,7 +201,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     #                    (test (comp_for | (',' test)* [','])) )
     for child in node.children:
       self.Visit(child)
-      if pytree_utils.NodeName(child) == 'COLON':
+      if child.type == grammar_token.COLON:
         # This is a key to a dictionary. We don't want to split the key if at
         # all possible.
         _SetStronglyConnected(child)
@@ -229,7 +230,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
           _SetSplitPenalty(
               pytree_utils.FirstLeafNode(node.children[1]),
               ONE_ELEMENT_ARGUMENT)
-      elif (pytree_utils.NodeName(node.children[0]) == 'LSQB' and
+      elif (node.children[0].type == grammar_token.LSQB and
             len(node.children[1].children) > 2 and
             (name.endswith('_test') or name.endswith('_expr'))):
         _SetStronglyConnected(node.children[1].children[0])
@@ -347,7 +348,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     _SetSplitPenalty(pytree_utils.FirstLeafNode(node), 0)
     prev_child = None
     for child in node.children:
-      if prev_child and pytree_utils.NodeName(prev_child) == 'COMMA':
+      if prev_child and prev_child.type == grammar_token.COMMA:
         _SetSplitPenalty(pytree_utils.FirstLeafNode(child), 0)
       prev_child = child
 
diff --git a/yapf/yapflib/subtype_assigner.py b/yapf/yapflib/subtype_assigner.py
index 1a54522..cb78809 100644
--- a/yapf/yapflib/subtype_assigner.py
+++ b/yapf/yapflib/subtype_assigner.py
@@ -26,7 +26,7 @@ Annotations:
 """
 
 from lib2to3 import pytree
-from lib2to3.pgen2 import token
+from lib2to3.pgen2 import token as grammar_token
 from lib2to3.pygram import python_symbols as syms
 
 from yapf.yapflib import format_token
@@ -75,14 +75,14 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
         comp_for = True
         _AppendFirstLeafTokenSubtype(child,
                                      format_token.Subtype.DICT_SET_GENERATOR)
-      elif pytree_utils.NodeName(child) in ('COLON', 'DOUBLESTAR'):
+      elif child.type in (grammar_token.COLON, grammar_token.DOUBLESTAR):
         dict_maker = True
 
     if not comp_for and dict_maker:
       last_was_colon = False
       unpacking = False
       for child in node.children:
-        if pytree_utils.NodeName(child) == 'DOUBLESTAR':
+        if child.type == grammar_token.DOUBLESTAR:
           _AppendFirstLeafTokenSubtype(child,
                                        format_token.Subtype.KWARGS_STAR_STAR)
         if last_was_colon:
@@ -100,8 +100,8 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
             _AppendFirstLeafTokenSubtype(child,
                                          format_token.Subtype.DICTIONARY_KEY)
           _AppendSubtypeRec(child, format_token.Subtype.DICTIONARY_KEY_PART)
-        last_was_colon = pytree_utils.NodeName(child) == 'COLON'
-        if pytree_utils.NodeName(child) == 'DOUBLESTAR':
+        last_was_colon = child.type == grammar_token.COLON
+        if child.type == grammar_token.DOUBLESTAR:
           unpacking = True
         elif last_was_colon:
           unpacking = False
@@ -274,7 +274,7 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
     # funcdef ::=
     #     'def' NAME parameters ['->' test] ':' suite
     for child in node.children:
-      if pytree_utils.NodeName(child) == 'NAME' and child.value != 'def':
+      if child.type == grammar_token.NAME and child.value != 'def':
         _AppendTokenSubtype(child, format_token.Subtype.FUNC_DEF)
         break
     for child in node.children:
@@ -311,10 +311,10 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
     for i in range(1, len(node.children)):
       prev_child = node.children[i - 1]
       child = node.children[i]
-      if pytree_utils.NodeName(prev_child) == 'COMMA':
+      if prev_child.type == grammar_token.COMMA:
         _AppendFirstLeafTokenSubtype(child,
                                      format_token.Subtype.PARAMETER_START)
-      elif pytree_utils.NodeName(child) == 'COMMA':
+      elif child.type == grammar_token.COMMA:
         _AppendLastLeafTokenSubtype(prev_child,
                                     format_token.Subtype.PARAMETER_STOP)
 
@@ -322,9 +322,9 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
         tname = True
         _SetArgListSubtype(child, format_token.Subtype.TYPED_NAME,
                            format_token.Subtype.TYPED_NAME_ARG_LIST)
-      elif pytree_utils.NodeName(child) == 'COMMA':
+      elif child.type == grammar_token.COMMA:
         tname = False
-      elif pytree_utils.NodeName(child) == 'EQUAL' and tname:
+      elif child.type == grammar_token.EQUAL and tname:
         _AppendTokenSubtype(child, subtype=format_token.Subtype.TYPED_NAME)
         tname = False
 
@@ -436,14 +436,14 @@ def _InsertPseudoParentheses(node):
   """Insert pseudo parentheses so that dicts can be formatted correctly."""
   comment_node = None
   if isinstance(node, pytree.Node):
-    if node.children[-1].type == token.COMMENT:
+    if node.children[-1].type == grammar_token.COMMENT:
       comment_node = node.children[-1].clone()
       node.children[-1].remove()
 
   first = pytree_utils.FirstLeafNode(node)
   last = pytree_utils.LastLeafNode(node)
 
-  if first == last and first.type == token.COMMENT:
+  if first == last and first.type == grammar_token.COMMENT:
     # A comment was inserted before the value, which is a pytree.Leaf.
     # Encompass the dictionary's value into an ATOM node.
     last = first.next_sibling
@@ -462,17 +462,19 @@ def _InsertPseudoParentheses(node):
     last = pytree_utils.LastLeafNode(node)
 
   lparen = pytree.Leaf(
-      token.LPAR, u'(', context=('', (first.get_lineno(), first.column - 1)))
+      grammar_token.LPAR,
+      u'(',
+      context=('', (first.get_lineno(), first.column - 1)))
   last_lineno = last.get_lineno()
-  if last.type == token.STRING and '\n' in last.value:
+  if last.type == grammar_token.STRING and '\n' in last.value:
     last_lineno += last.value.count('\n')
 
-  if last.type == token.STRING and '\n' in last.value:
+  if last.type == grammar_token.STRING and '\n' in last.value:
     last_column = len(last.value.split('\n')[-1]) + 1
   else:
     last_column = last.column + len(last.value) + 1
   rparen = pytree.Leaf(
-      token.RPAR, u')', context=('', (last_lineno, last_column)))
+      grammar_token.RPAR, u')', context=('', (last_lineno, last_column)))
 
   lparen.is_pseudo = True
   rparen.is_pseudo = True