diff options
Diffstat (limited to 'yapf/yapflib/split_penalty.py')
-rw-r--r-- | yapf/yapflib/split_penalty.py | 11 |
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 |