Refactor common function.

6 files changed, 65 insertions, 81 deletions

diff --git a/yapf/yapflib/blank_line_calculator.py b/yapf/yapflib/blank_line_calculator.py
index 0880a80..0a4fcc2 100644
--- a/yapf/yapflib/blank_line_calculator.py
+++ b/yapf/yapflib/blank_line_calculator.py
@@ -116,7 +116,7 @@ class _BlankLineCalculator(pytree_visitor.PyTreeVisitor):
     """
     if self.last_was_class_or_function:
       if pytree_utils.NodeName(node) in _PYTHON_STATEMENTS:
-        leaf = _GetFirstChildLeaf(node)
+        leaf = pytree_utils.FirstLeafNode(node)
         self._SetNumNewlines(leaf, self._GetNumNewlines(leaf))
     self.last_was_class_or_function = False
     super(_BlankLineCalculator, self).DefaultNodeVisit(node)
@@ -169,16 +169,10 @@ class _BlankLineCalculator(pytree_visitor.PyTreeVisitor):
 
 
 def _StartsInZerothColumn(node):
-  return (_GetFirstChildLeaf(node).column == 0 or
+  return (pytree_utils.FirstLeafNode(node).column == 0 or
           (_AsyncFunction(node) and node.prev_sibling.column == 0))
 
 
 def _AsyncFunction(node):
   return (py3compat.PY3 and node.prev_sibling and
           pytree_utils.NodeName(node.prev_sibling) == 'ASYNC')
-
-
-def _GetFirstChildLeaf(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _GetFirstChildLeaf(node.children[0])
diff --git a/yapf/yapflib/identify_container.py b/yapf/yapflib/identify_container.py
index 9d38239..23630c5 100644
--- a/yapf/yapflib/identify_container.py
+++ b/yapf/yapflib/identify_container.py
@@ -50,10 +50,10 @@ class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
     if pytree_utils.NodeName(node.children[1]) == 'arglist':
       for child in node.children[1].children:
         pytree_utils.SetOpeningBracket(
-            _GetFirstLeafNode(child), node.children[0])
+            pytree_utils.FirstLeafNode(child), node.children[0])
     else:
       pytree_utils.SetOpeningBracket(
-          _GetFirstLeafNode(node.children[1]), node.children[0])
+          pytree_utils.FirstLeafNode(node.children[1]), node.children[0])
 
   def Visit_atom(self, node):  # pylint: disable=invalid-name
     for child in node.children:
@@ -65,10 +65,5 @@ class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
       return
 
     for child in node.children[1].children:
-      pytree_utils.SetOpeningBracket(_GetFirstLeafNode(child), node.children[0])
-
-
-def _GetFirstLeafNode(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _GetFirstLeafNode(node.children[0])
+      pytree_utils.SetOpeningBracket(
+          pytree_utils.FirstLeafNode(child), node.children[0])
diff --git a/yapf/yapflib/pytree_unwrapper.py b/yapf/yapflib/pytree_unwrapper.py
index bbb5cdb..049dbbb 100644
--- a/yapf/yapflib/pytree_unwrapper.py
+++ b/yapf/yapflib/pytree_unwrapper.py
@@ -377,11 +377,6 @@ def _ContainsComments(node):
 
 def _SetMustSplitOnFirstLeaf(node):
   """Set the "must split" annotation on the first leaf node."""
-
-  def FindFirstLeaf(node):
-    if isinstance(node, pytree.Leaf):
-      return node
-    return FindFirstLeaf(node.children[0])
-
   pytree_utils.SetNodeAnnotation(
-      FindFirstLeaf(node), pytree_utils.Annotation.MUST_SPLIT, True)
+      pytree_utils.FirstLeafNode(node), pytree_utils.Annotation.MUST_SPLIT,
+      True)
diff --git a/yapf/yapflib/pytree_utils.py b/yapf/yapflib/pytree_utils.py
index d89588f..999ba88 100644
--- a/yapf/yapflib/pytree_utils.py
+++ b/yapf/yapflib/pytree_utils.py
@@ -69,6 +69,18 @@ def NodeName(node):
     return pygram.python_grammar.number2symbol[node.type]
 
 
+def FirstLeafNode(node):
+  if isinstance(node, pytree.Leaf):
+    return node
+  return FirstLeafNode(node.children[0])
+
+
+def LastLeafNode(node):
+  if isinstance(node, pytree.Leaf):
+    return node
+  return LastLeafNode(node.children[-1])
+
+
 # lib2to3 thoughtfully provides pygram.python_grammar_no_print_statement for
 # parsing Python 3 code that wouldn't parse otherwise (when 'print' is used in a
 # context where a keyword is disallowed).
diff --git a/yapf/yapflib/split_penalty.py b/yapf/yapflib/split_penalty.py
index fe32e49..416eda3 100644
--- a/yapf/yapflib/split_penalty.py
+++ b/yapf/yapflib/split_penalty.py
@@ -106,7 +106,8 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     _SetUnbreakable(node.children[colon_idx])
     self.DefaultNodeVisit(node)
     if arrow_idx > 0:
-      _SetSplitPenalty(_LastChildNode(node.children[arrow_idx - 1]), 0)
+      _SetSplitPenalty(
+          pytree_utils.LastLeafNode(node.children[arrow_idx - 1]), 0)
       _SetUnbreakable(node.children[arrow_idx])
       _SetStronglyConnected(node.children[arrow_idx + 1])
 
@@ -155,9 +156,10 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     while index < len(node.children) - 1:
       child = node.children[index]
       if isinstance(child, pytree.Leaf) and child.value == '=':
-        _SetSplitPenalty(_FirstChildNode(node.children[index]), NAMED_ASSIGN)
         _SetSplitPenalty(
-            _FirstChildNode(node.children[index + 1]), NAMED_ASSIGN)
+            pytree_utils.FirstLeafNode(node.children[index]), NAMED_ASSIGN)
+        _SetSplitPenalty(
+            pytree_utils.FirstLeafNode(node.children[index + 1]), NAMED_ASSIGN)
       index += 1
 
   def Visit_tname(self, node):  # pylint: disable=invalid-name
@@ -167,9 +169,10 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     while index < len(node.children) - 1:
       child = node.children[index]
       if isinstance(child, pytree.Leaf) and child.value == ':':
-        _SetSplitPenalty(_FirstChildNode(node.children[index]), NAMED_ASSIGN)
         _SetSplitPenalty(
-            _FirstChildNode(node.children[index + 1]), NAMED_ASSIGN)
+            pytree_utils.FirstLeafNode(node.children[index]), NAMED_ASSIGN)
+        _SetSplitPenalty(
+            pytree_utils.FirstLeafNode(node.children[index + 1]), NAMED_ASSIGN)
       index += 1
 
   def Visit_dotted_name(self, node):  # pylint: disable=invalid-name
@@ -203,10 +206,11 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
         if (len(node.children[1].children) > 1 and
             pytree_utils.NodeName(node.children[1].children[1]) == 'comp_for'):
           # Don't penalize splitting before a comp_for expression.
-          _SetSplitPenalty(_FirstChildNode(node.children[1]), 0)
+          _SetSplitPenalty(pytree_utils.FirstLeafNode(node.children[1]), 0)
         else:
           _SetSplitPenalty(
-              _FirstChildNode(node.children[1]), ONE_ELEMENT_ARGUMENT)
+              pytree_utils.FirstLeafNode(node.children[1]),
+              ONE_ELEMENT_ARGUMENT)
       elif (pytree_utils.NodeName(node.children[0]) == 'LSQB' and
             len(node.children[1].children) > 2 and
             (name.endswith('_test') or name.endswith('_expr'))):
@@ -219,9 +223,11 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
                         (name.endswith('_expr') and
                          style.Get('SPLIT_BEFORE_BITWISE_OPERATOR')))
         if split_before:
-          _SetSplitPenalty(_LastChildNode(node.children[1].children[1]), 0)
+          _SetSplitPenalty(
+              pytree_utils.LastLeafNode(node.children[1].children[1]), 0)
         else:
-          _SetSplitPenalty(_FirstChildNode(node.children[1].children[2]), 0)
+          _SetSplitPenalty(
+              pytree_utils.FirstLeafNode(node.children[1].children[2]), 0)
 
         # Don't split the ending bracket of a subscript list.
         _SetVeryStronglyConnected(node.children[-1])
@@ -285,13 +291,15 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
           subtypes = pytree_utils.GetNodeAnnotation(
               trailer.children[0], pytree_utils.Annotation.SUBTYPE)
           if subtypes and format_token.Subtype.SUBSCRIPT_BRACKET in subtypes:
-            _SetStronglyConnected(_FirstChildNode(trailer.children[1]))
+            _SetStronglyConnected(
+                pytree_utils.FirstLeafNode(trailer.children[1]))
 
-          last_child_node = _LastChildNode(trailer)
+          last_child_node = pytree_utils.LastLeafNode(trailer)
           if last_child_node.value.strip().startswith('#'):
             last_child_node = last_child_node.prev_sibling
           if not style.Get('DEDENT_CLOSING_BRACKETS'):
-            if _LastChildNode(last_child_node.prev_sibling).value != ',':
+            last = pytree_utils.LastLeafNode(last_child_node.prev_sibling)
+            if last.value != ',':
               if last_child_node.value == ']':
                 _SetUnbreakable(last_child_node)
               else:
@@ -313,7 +321,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
 
   def Visit_comp_for(self, node):  # pylint: disable=invalid-name
     # comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter]
-    _SetSplitPenalty(_FirstChildNode(node), 0)
+    _SetSplitPenalty(pytree_utils.FirstLeafNode(node), 0)
     _SetStronglyConnected(*node.children[1:])
     self.DefaultNodeVisit(node)
 
@@ -331,10 +339,11 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     index = 1
     while index + 1 < len(node.children):
       if style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR'):
-        _DecrementSplitPenalty(_FirstChildNode(node.children[index]), OR_TEST)
+        _DecrementSplitPenalty(
+            pytree_utils.FirstLeafNode(node.children[index]), OR_TEST)
       else:
         _DecrementSplitPenalty(
-            _FirstChildNode(node.children[index + 1]), OR_TEST)
+            pytree_utils.FirstLeafNode(node.children[index + 1]), OR_TEST)
       index += 2
 
   def Visit_and_test(self, node):  # pylint: disable=invalid-name
@@ -344,10 +353,11 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     index = 1
     while index + 1 < len(node.children):
       if style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR'):
-        _DecrementSplitPenalty(_FirstChildNode(node.children[index]), AND_TEST)
+        _DecrementSplitPenalty(
+            pytree_utils.FirstLeafNode(node.children[index]), AND_TEST)
       else:
         _DecrementSplitPenalty(
-            _FirstChildNode(node.children[index + 1]), AND_TEST)
+            pytree_utils.FirstLeafNode(node.children[index + 1]), AND_TEST)
       index += 2
 
   def Visit_not_test(self, node):  # pylint: disable=invalid-name
@@ -359,8 +369,10 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     # comparison ::= expr (comp_op expr)*
     self.DefaultNodeVisit(node)
     if len(node.children) == 3 and _StronglyConnectedCompOp(node):
-      _SetSplitPenalty(_FirstChildNode(node.children[1]), STRONGLY_CONNECTED)
-      _SetSplitPenalty(_FirstChildNode(node.children[2]), STRONGLY_CONNECTED)
+      _SetSplitPenalty(
+          pytree_utils.FirstLeafNode(node.children[1]), STRONGLY_CONNECTED)
+      _SetSplitPenalty(
+          pytree_utils.FirstLeafNode(node.children[2]), STRONGLY_CONNECTED)
     else:
       _IncreasePenalty(node, COMPARISON)
 
@@ -381,7 +393,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
           _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
         else:
           _SetSplitPenalty(
-              _FirstChildNode(node.children[index + 1]),
+              pytree_utils.FirstLeafNode(node.children[index + 1]),
               style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
       index += 1
 
@@ -411,7 +423,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     while index < len(node.children) - 1:
       child = node.children[index]
       if pytree_utils.NodeName(child) in self._ARITH_OPS:
-        next_node = _FirstChildNode(node.children[index + 1])
+        next_node = pytree_utils.FirstLeafNode(node.children[index + 1])
         _SetSplitPenalty(next_node, ARITH_EXPR)
       index += 1
 
@@ -426,7 +438,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
     while index < len(node.children) - 1:
       child = node.children[index]
       if pytree_utils.NodeName(child) in self._TERM_OPS:
-        next_node = _FirstChildNode(node.children[index + 1])
+        next_node = pytree_utils.FirstLeafNode(node.children[index + 1])
         _SetSplitPenalty(next_node, TERM)
       index += 1
 
@@ -446,7 +458,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
           _SetSplitPenalty(node.children[-1], STRONGLY_CONNECTED)
         else:
           if len(node.children) > 2:
-            _SetSplitPenalty(_FirstChildNode(node.children[1]), EXPR)
+            _SetSplitPenalty(pytree_utils.FirstLeafNode(node.children[1]), EXPR)
           _SetSplitPenalty(node.children[-1], ATOM)
     elif node.children[0].value in '[{' and len(node.children) == 2:
       # Keep empty containers together if we can.
@@ -461,7 +473,7 @@ class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
         prev_was_comma = True
       else:
         if prev_was_comma:
-          _SetSplitPenalty(_FirstChildNode(child), 0)
+          _SetSplitPenalty(pytree_utils.FirstLeafNode(child), 0)
         prev_was_comma = False
 
   ############################################################################
@@ -513,7 +525,7 @@ def _SetExpressionPenalty(node, penalty):
       for child in node.children:
         RecExpression(child, first_child_leaf)
 
-  RecExpression(node, _FirstChildNode(node))
+  RecExpression(node, pytree_utils.FirstLeafNode(node))
 
 
 def _IncreasePenalty(node, amt):
@@ -533,7 +545,7 @@ def _IncreasePenalty(node, amt):
       for child in node.children:
         RecExpression(child, first_child_leaf)
 
-  RecExpression(node, _FirstChildNode(node))
+  RecExpression(node, pytree_utils.FirstLeafNode(node))
 
 
 def _RecAnnotate(tree, annotate_name, annotate_value):
@@ -559,8 +571,8 @@ def _RecAnnotate(tree, annotate_name, annotate_value):
 def _StronglyConnectedCompOp(op):
   if (len(op.children[1].children) == 2 and
       pytree_utils.NodeName(op.children[1]) == 'comp_op' and
-      _FirstChildNode(op.children[1]).value == 'not' and
-      _LastChildNode(op.children[1]).value == 'in'):
+      pytree_utils.FirstLeafNode(op.children[1]).value == 'not' and
+      pytree_utils.LastLeafNode(op.children[1]).value == 'in'):
     return True
   if (isinstance(op.children[1], pytree.Leaf) and
       op.children[1].value in {'==', 'in'}):
@@ -598,15 +610,3 @@ def _AllowBuilderStyleCalls(node):
       if prev_child.lineno != child.lineno:
         _SetSplitPenalty(child, 0)
     prev_child = child
-
-
-def _FirstChildNode(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _FirstChildNode(node.children[0])
-
-
-def _LastChildNode(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _LastChildNode(node.children[-1])
diff --git a/yapf/yapflib/subtype_assigner.py b/yapf/yapflib/subtype_assigner.py
index 8d3b7f4..8bf3d8d 100644
--- a/yapf/yapflib/subtype_assigner.py
+++ b/yapf/yapflib/subtype_assigner.py
@@ -96,7 +96,7 @@ class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
           # Mark the first leaf of a key entry as a DICTIONARY_KEY. We
           # normally want to split before them if the dictionary cannot exist
           # on a single line.
-          if not unpacking or _GetFirstLeafNode(child).value == '**':
+          if not unpacking or pytree_utils.FirstLeafNode(child).value == '**':
             _AppendFirstLeafTokenSubtype(child,
                                          format_token.Subtype.DICTIONARY_KEY)
           _AppendSubtypeRec(child, format_token.Subtype.DICTIONARY_KEY_PART)
@@ -384,8 +384,8 @@ def _InsertPseudoParentheses(node):
       comment_node = node.children[-1].clone()
       node.children[-1].remove()
 
-  first = _GetFirstLeafNode(node)
-  last = _GetLastLeafNode(node)
+  first = pytree_utils.FirstLeafNode(node)
+  last = pytree_utils.LastLeafNode(node)
 
   if first == last and first.type == token.COMMENT:
     # A comment was inserted before the value, which is a pytree.Leaf.
@@ -396,8 +396,8 @@ def _InsertPseudoParentheses(node):
     node = new_node
     last.remove()
 
-    first = _GetFirstLeafNode(node)
-    last = _GetLastLeafNode(node)
+    first = pytree_utils.FirstLeafNode(node)
+    last = pytree_utils.LastLeafNode(node)
 
   lparen = pytree.Leaf(
       token.LPAR, u'(', context=('', (first.get_lineno(), first.column - 1)))
@@ -426,15 +426,3 @@ def _InsertPseudoParentheses(node):
     new_node = pytree.Node(syms.atom, [lparen, clone, rparen])
     node.replace(new_node)
     _AppendFirstLeafTokenSubtype(clone, format_token.Subtype.DICTIONARY_VALUE)
-
-
-def _GetFirstLeafNode(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _GetFirstLeafNode(node.children[0])
-
-
-def _GetLastLeafNode(node):
-  if isinstance(node, pytree.Leaf):
-    return node
-  return _GetLastLeafNode(node.children[-1])