Compare the node's type instead of its name
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 @@
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 @@
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 _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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
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 @@
# (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 @@
_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 @@
_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 @@
"""
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 @@
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 @@
_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 @@
# 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 @@
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 @@
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 @@
"""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 @@
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
