Annotate AST trees with source code positions
—
Quality
Pending
Does it follow best practices?
Impact
Pending
No eval scenarios have been run
Helper functions for working with AST nodes, including type checking, tree traversal, and node classification utilities. These functions provide support for both standard Python ast module and astroid library nodes.
Functions for walking and visiting AST trees with fine-grained control over traversal order and node inclusion.
# Available through asttokens.util module
from asttokens.util import (
walk, visit_tree, iter_children_ast, iter_children_astroid,
iter_children_func
)
def walk(node, include_joined_str=False) -> Iterator:
"""
Pre-order traversal of AST tree yielding all descendant nodes.
Parameters:
- node (ast.AST): Root node to start traversal from
- include_joined_str (bool): Include f-string nodes that may have broken positions
Yields:
ast.AST: Each node in pre-order traversal
"""
def visit_tree(node, previsit, postvisit):
"""
Depth-first tree traversal with explicit stack and pre/post visit callbacks.
Parameters:
- node (ast.AST): Root node to traverse
- previsit (callable): Function called before visiting children
- postvisit (callable): Function called after visiting children
"""
def iter_children_ast(node, include_joined_str=False) -> Iterator:
"""
Iterate over direct children of an ast module node.
Parameters:
- node (ast.AST): Node to get children from
- include_joined_str (bool): Include f-string nodes
Yields:
ast.AST: Each child node
"""
def iter_children_astroid(node, include_joined_str=False) -> Iterator:
"""
Iterate over direct children of an astroid node.
Parameters:
- node: Astroid node to get children from
- include_joined_str (bool): Include f-string nodes
Yields:
Astroid node: Each child node
"""
def iter_children_func(node) -> callable:
"""
Return appropriate child iterator for AST or Astroid nodes.
Parameters:
- node: AST or Astroid node
Returns:
callable: iter_children_ast or iter_children_astroid function
"""import asttokens
import ast
source = """
def greet(name):
if name:
return f"Hello, {name}!"
return "Hello, World!"
"""
tree = ast.parse(source)
# Walk all nodes in tree
for node in asttokens.util.walk(tree):
print(f"{type(node).__name__}: {getattr(node, 'lineno', 'N/A')}")
# Custom tree traversal with callbacks
def pre_visit(node, parent, field):
print(f"Entering: {type(node).__name__}")
def post_visit(node, parent, field):
print(f"Exiting: {type(node).__name__}")
asttokens.util.visit_tree(tree, pre_visit, post_visit)
# Iterate children of function node
func_node = tree.body[0]
for child in asttokens.util.iter_children_ast(func_node):
print(f"Child: {type(child).__name__}")Functions for identifying different categories of AST nodes by their semantic role.
# Available through asttokens.util module
from asttokens.util import (
is_expr, is_stmt, is_module, is_joined_str, is_expr_stmt,
is_constant, is_ellipsis, is_starred, is_slice, is_empty_astroid_slice
)
def is_expr(node) -> bool:
"""
Check if node is an expression.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is an expression
"""
def is_stmt(node) -> bool:
"""
Check if node is a statement.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is a statement
"""
def is_module(node) -> bool:
"""
Check if node is a module.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is a Module
"""
def is_joined_str(node) -> bool:
"""
Check if node is an f-string (JoinedStr).
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is a JoinedStr (f-string)
"""
def is_expr_stmt(node) -> bool:
"""
Check if node is an expression statement.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is an Expr statement
"""
def is_constant(node) -> bool:
"""
Check if node represents a constant value.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is a constant (ast.Constant or astroid.Const)
"""
def is_ellipsis(node) -> bool:
"""
Check if node represents ellipsis (...).
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node represents ellipsis
"""
def is_starred(node) -> bool:
"""
Check if node is a starred expression (*args).
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node is a Starred expression
"""
def is_slice(node) -> bool:
"""
Check if node represents a slice operation.
Parameters:
- node (ast.AST): Node to check
Returns:
bool: True if node represents slice
"""
def is_empty_astroid_slice(node) -> bool:
"""
Check if node is an empty Astroid slice.
Parameters:
- node: Node to check (Astroid)
Returns:
bool: True if node is empty Astroid slice
"""import asttokens
import ast
source = """
x = 42
print(x)
y = [1, 2, 3]
z = y[1:3]
def func(*args): pass
"""
tree = ast.parse(source)
# Check node types
for node in ast.walk(tree):
if asttokens.util.is_stmt(node):
print(f"Statement: {type(node).__name__}")
elif asttokens.util.is_expr(node):
print(f"Expression: {type(node).__name__}")
if asttokens.util.is_constant(node):
print(f"Constant: {node.value}")
if asttokens.util.is_starred(node):
print("Found starred expression")
if asttokens.util.is_slice(node):
print("Found slice operation")
# Check specific patterns
assign_node = tree.body[0] # x = 42
print(f"Is statement: {asttokens.util.is_stmt(assign_node)}")
constant_node = assign_node.value # 42
print(f"Is constant: {asttokens.util.is_constant(constant_node)}")Functions for working with specific types of statements and expressions.
# Available through asttokens.util module
from asttokens.util import (
last_stmt, replace, annotate_fstring_nodes, fstring_positions_work
)
def last_stmt(node) -> ast.AST:
"""
Get the last statement in a multi-statement node.
Parameters:
- node (ast.AST): Node that may contain multiple statements
Returns:
ast.AST: Last statement node
"""
def replace(text, replacements) -> str:
"""
Apply multiple text replacements to source code.
Parameters:
- text (str): Original text
- replacements (List[Tuple[int, int, str]]): List of (start, end, replacement) tuples
Returns:
str: Text with replacements applied
"""
def annotate_fstring_nodes(tree):
"""
Mark f-string nodes that have broken position information.
Parameters:
- tree (ast.AST): AST tree to annotate
"""
def fstring_positions_work() -> bool:
"""
Check if f-string positions work correctly in current Python version.
Returns:
bool: True if f-string positions are reliable
"""import asttokens
import ast
source = """
if True:
print("first")
print("second")
x = 42
"""
tree = ast.parse(source)
if_node = tree.body[0]
# Get last statement in if block
last = asttokens.util.last_stmt(if_node.body[0])
print(f"Last statement: {type(last).__name__}")
# Check f-string support
if asttokens.util.fstring_positions_work():
print("F-string positions work correctly")
# Apply text replacements
original = "hello world"
replacements = [(0, 5, "hi"), (6, 11, "there")]
result = asttokens.util.replace(original, replacements)
print(result) # "hi there"
# Annotate f-strings in tree (modifies tree in-place)
asttokens.util.annotate_fstring_nodes(tree)Additional utility functions for working with tokens and text manipulation.
# Available through asttokens.util module
from asttokens.util import combine_tokens
def combine_tokens(group) -> list:
"""
Combine consecutive NAME, NUMBER, and ERRORTOKEN tokens.
Parameters:
- group (List[tokenize.TokenInfo]): List of tokens to combine
Returns:
List[tokenize.TokenInfo]: Combined token list
"""import asttokens
import asttokens.util
# The combine_tokens function is used internally for token processing
# but is available for advanced token manipulation scenarios
source = "hello_world123"
tokens = list(asttokens.util.generate_tokens(source))
print([t.string for t in tokens])Install with Tessl CLI
npx tessl i tessl/pypi-asttokens