Merge pull request #116 from gristlabs/walk-include-joinedstr

Add include_joined_str parameter to util.walk

asttokens/util.py

@@ -133,10 +133,9 @@ def iter_children_func(node):
   return iter_children_astroid if hasattr(node, 'get_children') else iter_children_ast
 
 
-def iter_children_astroid(node):
-  # type: (NodeNG) -> Union[Iterator, List]
-  # Don't attempt to process children of JoinedStr nodes, which we can't fully handle yet.
-  if is_joined_str(node):
+def iter_children_astroid(node, include_joined_str=False):
+  # type: (NodeNG, bool) -> Union[Iterator, List]
+  if not include_joined_str and is_joined_str(node):
     return []
 
   return node.get_children()
@@ -145,10 +144,10 @@ def iter_children_astroid(node):
 SINGLETONS = {c for n, c in iteritems(ast.__dict__) if isinstance(c, type) and
               issubclass(c, (ast.expr_context, ast.boolop, ast.operator, ast.unaryop, ast.cmpop))}
 
-def iter_children_ast(node):
-  # type: (AST) -> Iterator[Union[AST, expr]]
-  # Don't attempt to process children of JoinedStr nodes, which we can't fully handle yet.
-  if is_joined_str(node):
+
+def iter_children_ast(node, include_joined_str=False):
+  # type: (AST, bool) -> Iterator[Union[AST, expr]]
+  if not include_joined_str and is_joined_str(node):
     return
 
   if isinstance(node, ast.Dict):
@@ -274,15 +273,17 @@ def visit_tree(node, previsit, postvisit):
   return ret
 
 
-
-def walk(node):
-  # type: (AST) -> Iterator[Union[Module, AstNode]]
+def walk(node, include_joined_str=False):
+  # type: (AST, bool) -> Iterator[Union[Module, AstNode]]
   """
   Recursively yield all descendant nodes in the tree starting at ``node`` (including ``node``
   itself), using depth-first pre-order traversal (yieling parents before their children).
 
   This is similar to ``ast.walk()``, but with a different order, and it works for both ``ast`` and
   ``astroid`` trees. Also, as ``iter_children()``, it skips singleton nodes generated by ``ast``.
+
+  By default, ``JoinedStr`` (f-string) nodes and their contents are skipped
+  because they previously couldn't be handled. Set ``include_joined_str`` to True to include them.
   """
   iter_children = iter_children_func(node)
   done = set()
@@ -297,7 +298,7 @@ def walk(node):
     # Insert all children in reverse order (so that first child ends up on top of the stack).
     # This is faster than building a list and reversing it.
     ins = len(stack)
-    for c in iter_children(current):
+    for c in iter_children(current, include_joined_str):
       stack.insert(ins, c)
 
 
@@ -457,7 +458,7 @@ def annotate_fstring_nodes(tree):
       # f-strings were weirdly implemented until https://peps.python.org/pep-0701/
       # In Python 3.12, inner nodes have sensible positions.
       return
-    for joinedstr in walk(tree):
+    for joinedstr in walk(tree, include_joined_str=True):
       if not isinstance(joinedstr, ast.JoinedStr):
         continue
       for part in joinedstr.values:
@@ -468,19 +469,11 @@ def annotate_fstring_nodes(tree):
           if not fstring_positions_work():
             for child in walk(part.value):
               setattr(child, '_broken_positions', True)
-              if isinstance(child, ast.JoinedStr):
-                # Recursively handle this inner JoinedStr in the same way.
-                # While this is usually automatic for other nodes,
-                # the children of f-strings are explicitly excluded in iter_children_ast.
-                annotate_fstring_nodes(child)
 
           if part.format_spec:  # this is another JoinedStr
             # Again, the standard positions span the full f-string.
             setattr(part.format_spec, '_broken_positions', True)
-            # Recursively handle this inner JoinedStr in the same way.
-            # While this is usually automatic for other nodes,
-            # the children of f-strings are explicitly excluded in iter_children_ast.
-            annotate_fstring_nodes(part.format_spec)
+
 else:
   def fstring_positions_work():
     # type: () -> bool