aboutsummaryrefslogtreecommitdiff
path: root/yapf/yapflib/split_penalty.py
diff options
context:
space:
mode:
Diffstat (limited to 'yapf/yapflib/split_penalty.py')
-rw-r--r--yapf/yapflib/split_penalty.py11
1 files changed, 6 insertions, 5 deletions
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