"""Prefix-preserving editing primitives for parso trees.
parso stores leading whitespace and comments on each leaf's ``prefix`` string;
the prefix of a non-leaf node is the prefix of its first leaf. Edits that
splice nodes in or out of the tree must transfer prefixes carefully so that
trailing comments stay glued to the right line, blank-line spacing does not
drift, and indent depth is preserved.
The primitives here operate directly on the parso tree (mutating
``parent.children`` lists). They are deliberately small — composition lives
in higher tiers (`scene_file`, `scene_module`).
"""
from __future__ import annotations
from parso.tree import Leaf, NodeOrLeaf
from .source_tree import parse_snippet
[docs]
def replace_node(old: NodeOrLeaf, new: NodeOrLeaf, *, preserve_prefix: bool = True) -> None:
"""Replace ``old`` with ``new`` in their shared parent's ``children`` list.
With ``preserve_prefix=True`` (default), the leading whitespace + comments
of ``old``'s first leaf are transferred onto ``new``'s first leaf so that
same-line trailing comments on the previous statement, leading blank
lines, and decorators above remain attached.
"""
parent = _require_parent(old, "replace_node")
if preserve_prefix:
_set_prefix(new, _get_prefix(old))
_swap_child(parent, old, new)
[docs]
def insert_after(sibling: NodeOrLeaf, new_node: NodeOrLeaf, *, copy_indent: bool = True) -> None:
"""Insert ``new_node`` immediately after ``sibling`` in their shared parent.
With ``copy_indent=True`` (default), the indent run of ``sibling``'s first
leaf prefix is copied onto ``new_node`` so the new statement sits at the
same column. ``new_node``'s prefix is overwritten with ``"\\n<indent>"`` —
callers wanting custom prefixes should pass ``copy_indent=False`` and
populate the prefix themselves.
"""
parent = _require_parent(sibling, "insert_after")
children = parent.children
idx = _index_of(children, sibling)
if copy_indent:
indent = _indent_of(_get_prefix(sibling))
_set_prefix(new_node, "\n" + indent)
new_node.parent = parent
children.insert(idx + 1, new_node)
[docs]
def insert_before(sibling: NodeOrLeaf, new_node: NodeOrLeaf, *, copy_indent: bool = True) -> None:
"""Insert ``new_node`` immediately before ``sibling`` in their shared parent.
With ``copy_indent=True`` (default), ``new_node`` inherits ``sibling``'s
full prefix (so any leading comments/blank lines stay above the inserted
node) and ``sibling``'s prefix is reset to ``"\\n<indent>"`` so it sits
at the same column it did originally.
Note: this transfers comments above ``sibling`` *to the inserted node*. To
keep them attached to ``sibling``, pass ``copy_indent=False`` and manage
prefixes manually.
"""
parent = _require_parent(sibling, "insert_before")
children = parent.children
idx = _index_of(children, sibling)
if copy_indent:
original_prefix = _get_prefix(sibling)
indent = _indent_of(original_prefix)
_set_prefix(new_node, original_prefix)
_set_prefix(sibling, "\n" + indent)
new_node.parent = parent
children.insert(idx, new_node)
[docs]
def remove_node(node: NodeOrLeaf, *, collapse_blank_lines: bool = True) -> None:
"""Remove ``node`` from its parent's ``children`` list.
With ``collapse_blank_lines=True`` (default), surplus blank lines in
``node``'s prefix are collapsed onto the next sibling so deleting
statements in sequence does not balloon vertical spacing. The collapse
rule is: keep at most **one** blank line of separation; the indent run on
the final line is preserved verbatim.
Same-line trailing comments stored in ``node``'s prefix (which originate
on the *previous* sibling — see module docstring) are re-attached to the
next sibling so they stay on their original line.
"""
parent = _require_parent(node, "remove_node")
children = parent.children
idx = _index_of(children, node)
next_sibling = children[idx + 1] if idx + 1 < len(children) else None
removed_prefix = _get_prefix(node)
children.pop(idx)
node.parent = None
if next_sibling is None:
return
next_prefix = _get_prefix(next_sibling)
merged = _merge_prefix_on_remove(removed_prefix, next_prefix, collapse_blank_lines=collapse_blank_lines)
_set_prefix(next_sibling, merged)
[docs]
def get_call_kwarg(call_node: NodeOrLeaf, name: str) -> NodeOrLeaf | None:
"""Return the value subtree for kwarg ``name`` in a call, or ``None``.
``call_node`` may be either the ``trailer`` (the ``(...)`` after a name)
or the enclosing ``atom_expr`` — both forms are accepted.
"""
trailer = _resolve_trailer(call_node)
if trailer is None:
return None
for arg in _iter_arguments(trailer):
if _argument_name(arg) == name:
return _argument_value(arg)
return None
[docs]
def set_call_kwarg(call_node: NodeOrLeaf, name: str, value_expr: str) -> None:
"""Set kwarg ``name`` on a call expression.
Overwrites if ``name`` already exists (preserves the order of other args);
appends if not. ``value_expr`` is parsed with :func:`parse_snippet`, so
callers pass real Python source (e.g. ``"Vec2(0, 0)"`` or ``'"hello"'``).
A trailing comma in the original ``arglist`` is preserved when appending.
"""
trailer = _resolve_trailer(call_node)
if trailer is None:
raise ValueError("set_call_kwarg: node is not a callable trailer")
new_value = parse_snippet(value_expr)
new_value.parent = None
for arg in _iter_arguments(trailer):
if _argument_name(arg) == name:
old_value = _argument_value(arg)
if old_value is None:
raise ValueError(f"set_call_kwarg: argument {name!r} has no value to replace")
_set_prefix(new_value, _get_prefix(old_value))
_swap_child(arg, old_value, new_value)
return
_append_kwarg(trailer, name, new_value)
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _require_parent(node: NodeOrLeaf, op: str) -> NodeOrLeaf:
parent = node.parent
if parent is None:
raise ValueError(f"{op}: node has no parent")
return parent
def _index_of(children: list[NodeOrLeaf], target: NodeOrLeaf) -> int:
for i, c in enumerate(children):
if c is target:
return i
raise ValueError("node not found in parent.children")
def _swap_child(parent: NodeOrLeaf, old: NodeOrLeaf, new: NodeOrLeaf) -> None:
children = parent.children
idx = _index_of(children, old)
new.parent = parent
old.parent = None
children[idx] = new
def _get_prefix(node: NodeOrLeaf) -> str:
if isinstance(node, Leaf):
return node.prefix
return node.get_first_leaf().prefix
def _set_prefix(node: NodeOrLeaf, prefix: str) -> None:
if isinstance(node, Leaf):
node.prefix = prefix
else:
node.get_first_leaf().prefix = prefix
def _indent_of(prefix: str) -> str:
"""Indent run = trailing run of spaces/tabs after the last newline.
For ``"\\n\\n "`` returns ``" "``; for ``" "`` returns the same;
for ``""`` returns ``""``.
"""
if "\n" in prefix:
tail = prefix.rsplit("\n", 1)[1]
else:
tail = prefix
# Tail may itself contain a comment-then-indent; the indent we want is the
# final whitespace run, so strip back to whitespace.
out = []
for ch in reversed(tail):
if ch in " \t":
out.append(ch)
else:
break
return "".join(reversed(out))
def _merge_prefix_on_remove(removed: str, next_prefix: str, *, collapse_blank_lines: bool) -> str:
"""Combine the removed node's prefix with the next sibling's prefix.
Strategy:
* Split each prefix into (comment_lines, indent_run).
* Comments before the removed node belonged to the previous sibling's
end-of-line and must be preserved verbatim if they ended *before* a
newline (ie they live on the previous source line). We capture them by
keeping any comment-bearing lines from ``removed`` that appear before
the first ``\\n``.
* Blank lines (purely-whitespace lines) in either prefix are collapsed
to at most one when ``collapse_blank_lines`` is true; otherwise both
prefixes' blank lines are concatenated as-is.
* The indent run of ``next_prefix`` (its trailing post-final-newline
whitespace) is preserved verbatim — the next sibling must keep its
column.
"""
if not collapse_blank_lines:
return removed + next_prefix
leading_comment, removed_rest = _split_inline_comment(removed)
# Fast path: when neither prefix carries a newline (typical for siblings
# inside an indented suite where each simple_stmt's prefix is just the
# indent run), the previous statement's trailing newline is what
# separates the lines. Concatenating ``removed_rest`` and ``next_prefix``
# would double-count the indent. Keep ``next_prefix`` verbatim.
if "\n" not in removed_rest and "\n" not in next_prefix:
return leading_comment + next_prefix
# Compute blank-line counts from each prefix separately. Concatenating
# ``removed_rest + next_prefix`` and reading its trailing indent run
# would double-count the suite indent — both prefixes typically carry
# the suite indent, and the combined trailing run is the sum of those
# two indents. The next sibling's own indent is the authoritative one.
removed_blanks, _ = _split_blank_lines_and_indent(removed_rest)
next_blanks, indent = _split_blank_lines_and_indent(next_prefix)
blank_lines = removed_blanks + next_blanks
if blank_lines > 1:
blank_lines = 1
out = leading_comment
if blank_lines == 0:
# Need at least one newline to terminate the previous statement when
# the next sibling exists; if leading_comment is empty and indent is
# also empty, we end up with an empty prefix which is correct (no
# newline needed in expression contexts like arglists).
if indent or out:
out += "\n"
out += indent
else:
out += "\n" * blank_lines
out += indent
return out
def _split_inline_comment(prefix: str) -> tuple[str, str]:
"""Pull off a comment that appears before the first newline.
Such a comment is a "same-line trailing comment" on the previous sibling.
Returns ``(comment_with_trailing_newline_or_empty, rest_of_prefix)``.
"""
if not prefix or "\n" not in prefix:
# Whole prefix is on a single line. A leading comment here belongs to
# the previous statement — keep it; otherwise nothing to extract.
if "#" in prefix:
return prefix + "\n", ""
return "", prefix
head, tail = prefix.split("\n", 1)
if "#" in head:
return head + "\n", tail
return "", prefix
def _split_blank_lines_and_indent(prefix: str) -> tuple[int, str]:
"""Count newlines in ``prefix`` (proxy for blank-line count) and return
the trailing indent run."""
if "\n" not in prefix:
return 0, prefix
lines = prefix.split("\n")
indent = lines[-1]
# The number of blank lines is the number of \n separators between
# whitespace/empty lines; for typical fixtures this matches the count of
# \n minus 0 (every \n marks the start of a new line; the last line is
# the indent of the next statement).
return len(lines) - 1, indent
def _resolve_trailer(node: NodeOrLeaf) -> NodeOrLeaf | None:
"""Return the call ``trailer`` for ``node`` if it represents one, else None.
Accepts either a ``trailer`` directly or an ``atom_expr`` whose final
child is a call trailer ``( ... )``.
"""
if node.type == "trailer" and _is_call_trailer(node):
return node
if node.type == "atom_expr":
last = node.children[-1]
if last.type == "trailer" and _is_call_trailer(last):
return last
return None
def _is_call_trailer(trailer: NodeOrLeaf) -> bool:
if not hasattr(trailer, "children") or not trailer.children:
return False
first = trailer.children[0]
return getattr(first, "type", None) == "operator" and first.value == "("
def _iter_arguments(trailer: NodeOrLeaf):
"""Yield ``argument`` nodes inside a call trailer."""
inner = trailer.children[1:-1] # strip ( and )
if not inner:
return
if len(inner) == 1:
node = inner[0]
if node.type == "argument":
yield node
elif node.type == "arglist":
for child in node.children:
if child.type == "argument":
yield child
return
# Trailer with multiple inner children only happens for arglists in
# practice; fall through to the same logic.
for child in inner:
if child.type == "argument":
yield child
def _argument_name(arg: NodeOrLeaf) -> str | None:
"""Return the kwarg name of an ``argument`` node, or None for positionals."""
if not hasattr(arg, "children") or len(arg.children) < 3:
return None
name_node, eq, _value = arg.children[0], arg.children[1], arg.children[2]
if name_node.type != "name" or getattr(eq, "type", None) != "operator" or eq.value != "=":
return None
return name_node.value
def _argument_value(arg: NodeOrLeaf) -> NodeOrLeaf | None:
if not hasattr(arg, "children") or len(arg.children) < 3:
return None
return arg.children[2]
def _append_kwarg(trailer: NodeOrLeaf, name: str, value: NodeOrLeaf) -> None:
"""Append ``name=value`` to a call's arglist, preserving trailing comma."""
children = trailer.children
open_paren = children[0]
close_paren = children[-1]
inner = children[1:-1]
new_arg = _build_kwarg_argument(name, value)
if not inner:
# Empty call: f() -> f(name=value)
new_arg.parent = trailer
children.insert(-1, new_arg)
return
if len(inner) == 1 and inner[0].type == "argument":
# Single-arg call: promote to arglist.
existing = inner[0]
comma = _make_op(",", prefix="")
_set_prefix(new_arg, " ")
arglist = _make_arglist([existing, comma, new_arg])
arglist.parent = trailer
children[1:-1] = [arglist]
# Preserve close-paren prefix unchanged.
_ = open_paren, close_paren
return
# Existing arglist case.
if len(inner) == 1 and inner[0].type == "arglist":
arglist = inner[0]
had_trailing_comma = arglist.children[-1].type == "operator" and arglist.children[-1].value == ","
if had_trailing_comma:
# Original layout: [..., last_arg, trailing_comma]. We want
# [..., last_arg, separator_comma, new_arg, trailing_comma]. Insert
# ``separator_comma`` then ``new_arg`` before the trailing comma.
new_arg.parent = arglist
_set_prefix(new_arg, " ")
separator = _make_op(",", prefix="")
separator.parent = arglist
arglist.children.insert(-1, separator)
arglist.children.insert(-1, new_arg)
else:
comma = _make_op(",", prefix="")
_set_prefix(new_arg, " ")
new_arg.parent = arglist
comma.parent = arglist
arglist.children.append(comma)
arglist.children.append(new_arg)
return
# Fallback: build a fresh arglist from whatever inner had.
raise ValueError(f"_append_kwarg: unexpected trailer inner shape {[c.type for c in inner]}")
def _build_kwarg_argument(name: str, value: NodeOrLeaf) -> NodeOrLeaf:
"""Construct a fresh ``argument`` parso node ``name=value`` from a value
snippet."""
snippet = parse_snippet(f"_({name}={_render_for_snippet(value)})")
# snippet is the atom_expr; descend to the argument node.
trailer = snippet.children[-1] # type: ignore[union-attr]
inner = trailer.children[1:-1]
if len(inner) == 1 and inner[0].type == "argument":
arg = inner[0]
else:
raise ValueError("_build_kwarg_argument: failed to lift argument node")
arg.parent = None
# Replace the parsed value (which was rendered via get_code) with the
# caller's actual subtree to preserve any sub-structure exactly.
parsed_value = arg.children[2]
value.parent = arg
arg.children[2] = value
_set_prefix(value, _get_prefix(parsed_value))
return arg
def _render_for_snippet(value: NodeOrLeaf) -> str:
"""Serialise ``value`` for round-tripping through :func:`parse_snippet`.
The serialised form is only used to build a syntactically valid argument
node; the real subtree replaces the rendered placeholder afterwards. We
strip the leading prefix to avoid spaces leaking into the wrapper call.
"""
code = value.get_code()
return code.lstrip()
def _make_op(value: str, *, prefix: str) -> Leaf:
"""Create a synthetic ``operator`` leaf via parso, attaching ``prefix``."""
# For commas, lift the comma leaf out of an arglist.
snippet = parse_snippet(f"f(a{value}b)") if value == "," else parse_snippet(f"x {value} y")
leaf = _find_op_leaf(snippet, value)
if leaf is None:
raise ValueError(f"_make_op: could not synthesise {value!r}")
leaf.parent = None
leaf.prefix = prefix
return leaf
def _find_op_leaf(node: NodeOrLeaf, value: str) -> Leaf | None:
if isinstance(node, Leaf):
if node.type == "operator" and node.value == value:
return node
return None
for c in node.children: # type: ignore[union-attr]
found = _find_op_leaf(c, value)
if found is not None:
return found
return None
def _make_arglist(args: list[NodeOrLeaf]) -> NodeOrLeaf:
"""Construct an ``arglist`` parso node from a list of ``argument``+``,``."""
snippet = parse_snippet("f(a, b)")
trailer = snippet.children[-1] # type: ignore[union-attr]
arglist = trailer.children[1]
if arglist.type != "arglist":
raise ValueError("_make_arglist: failed to lift arglist node")
arglist.children = list(args)
for c in args:
c.parent = arglist
arglist.parent = None
return arglist