"""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 collections.abc import Iterator
from parso.tree import BaseNode, Leaf, NodeOrLeaf
from .layout import INDENT, LINE_LIMIT, call_layout, carries_magic_comma, display_width, trailing_comma_is_syntax
from .source_tree import _children, is_leaf, is_node, 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 remove_statement(stmt: NodeOrLeaf, *, collapse_blank_lines: bool = True) -> None:
"""Remove one statement, which a semicolon may have joined to others on its line.
``stmt`` is a statement as parso hands one out: the sole statement of an
ordinary line, or one of the several a semicolon-joined line holds. A line
down to its last statement goes whole, newline and all
(:func:`remove_node`, whose prefix handling this then inherits); otherwise
only the statement goes, together with one of the semicolons beside it, and
the line keeps its indent whichever end it lost.
"""
line = stmt if stmt.type == "simple_stmt" else stmt.parent
if line is None or line.type != "simple_stmt" or not is_node(line):
remove_node(stmt, collapse_blank_lines=collapse_blank_lines)
return
children = _children(line)
if stmt is line or len(line_statements(line)) <= 1:
remove_node(line, collapse_blank_lines=collapse_blank_lines)
return
idx = _index_of(children, stmt)
prefix = _get_prefix(stmt)
if idx + 1 < len(children) and _is_semicolon(children[idx + 1]):
del children[idx : idx + 2]
if idx == 0:
# The line's indent lived on what has just gone; whatever now leads
# the line has to carry it instead.
_set_prefix(children[idx], prefix)
else:
del children[idx - 1 : idx + 1]
stmt.parent = None
[docs]
def line_statements(line: NodeOrLeaf) -> list[NodeOrLeaf]:
"""The statements written on one line, which a semicolon can join several of.
parso folds ``self.name = "Root"; self.add_child(Hero())`` into a single
``simple_stmt`` whose children are the two statements, the semicolon between
them and the newline, so a reader that looks only at the first child sees
half the line. Anything that is not a ``simple_stmt`` is already a single
statement and comes back on its own.
"""
if line.type != "simple_stmt" or not is_node(line):
return [line]
return [c for c in _children(line) if c.type != "newline" and not _is_semicolon(c)]
def _is_semicolon(node: NodeOrLeaf) -> bool:
return node.type == "operator" and getattr(node, "value", None) == ";"
[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)
break
else:
_append_kwarg(trailer, name, new_value)
relayout_statement(trailer)
[docs]
def set_call_positional(call_node: NodeOrLeaf, index: int, value_expr: str) -> None:
"""Replace the ``index``-th argument a call passes by position.
The counterpart of :func:`set_call_kwarg` for the arguments a call does not
name. Which parameter a position fills is a fact about the callable rather
than about the text, so establishing that ``index`` is the caller's
business; all that happens here is that the expression standing there is
swapped for ``value_expr`` and the statement laid out again.
Raises :class:`ValueError` when the call passes no such position.
"""
trailer = _resolve_trailer(call_node)
if trailer is None:
raise ValueError("set_call_positional: node is not a callable trailer")
positions = [item for item in iter_call_items(trailer) if not is_named_argument(item)]
if not 0 <= index < len(positions):
raise ValueError(f"set_call_positional: the call passes no argument at position {index}")
old = positions[index]
new_value = parse_snippet(value_expr)
new_value.parent = None
_set_prefix(new_value, _get_prefix(old))
_swap_child(_require_parent(old, "set_call_positional"), old, new_value)
relayout_statement(trailer)
[docs]
def iter_call_items(trailer: BaseNode) -> Iterator[NodeOrLeaf]:
"""Yield everything a call passes, named or not, in the order it is written.
:func:`_iter_arguments` yields the ``argument`` nodes, which is to say the
named ones and the unpackings; this yields those *and* the bare expressions
a call passes by position, so a caller counting positions counts them
against the same list the interpreter would.
"""
inner = _children(trailer)[1:-1]
if not inner:
return
if len(inner) == 1 and inner[0].type == "arglist" and is_node(inner[0]):
items: list[NodeOrLeaf] = _children(inner[0])
else:
items = list(inner)
for item in items:
if item.type == "operator" and getattr(item, "value", None) == ",":
continue
yield item
[docs]
def is_named_argument(item: NodeOrLeaf) -> bool:
"""Is this call item written as ``name=value``?
``False`` for everything else :func:`iter_call_items` yields: an expression
passed by position, and an unpacking (``*args``, ``**kwargs``), which parso
also builds an ``argument`` node for but which names no parameter.
"""
return is_node(item) and item.type == "argument" and _argument_name(item) is not None
[docs]
def relayout_statement(node: NodeOrLeaf, *, limit: int = LINE_LIMIT) -> None:
"""Lay the statement containing ``node`` out to fit ``limit`` columns.
Called after a call in an existing file has been edited, so the statement
the save just rewrote is left in the shape a formatter would leave it: on
one line while it fits, and one keyword argument per line once it does not
(:mod:`~simvx.core.scene_io.layout` decides, from the same rule the emitter
writes new files by). A statement nothing wrote to is never reached, so a
save with no edits in it changes no line's shape.
The author's own expressions survive it: nothing is reparsed, reprinted or
rewritten, and every value is the node it always was, down to its quotes and
its digits. What moves is the whitespace the layout owns and only that --
after an opening bracket, around a comma, before a closing bracket -- and it
moves at every depth, so a value the author spaced out inside its own
brackets comes back spaced the way a formatter spaces it.
Three statements are left exactly as found, because relaying them out would
destroy something this layer cannot put back: one carrying a comment, which
belongs to the line it was written on; one holding a value spanning
lines of its own (a triple-quoted string, a continuation); and one sharing
its line with other statements behind semicolons, whose width is not one
statement's to measure.
"""
stmt = enclosing_statement(node)
if stmt is None:
return
expression = _children(stmt)[0]
group = _trailing_group(expression)
if group is None or not _can_relayout(stmt, group):
return
_flatten(group)
indent = len(_indent_of(_get_prefix(stmt)))
prefix_width, suffix_width = _measure(list(_iter_leaves(stmt)), group, indent)
_lay_out_group(group, indent=indent, prefix_width=prefix_width, suffix_width=suffix_width, limit=limit)
[docs]
def enclosing_statement(node: NodeOrLeaf) -> BaseNode | None:
"""The ``simple_stmt`` ``node`` sits in, or ``None`` when it sits in none.
Worth asking for *before* an edit that detaches ``node`` from the tree: the
answer is what :func:`relayout_statement` then has to be given, since the
detached node no longer leads anywhere.
"""
current: NodeOrLeaf | None = node
while current is not None and current.type != "simple_stmt":
current = current.parent
return current if current is not None and is_node(current) else None
# ---------------------------------------------------------------------------
# Statement layout
# ---------------------------------------------------------------------------
#: Node types parso folds a bracket's comma-separated contents into. A bracket
#: holding one item has no such node: the item sits under the bracket directly.
_SEQUENCES = frozenset(
{"arglist", "testlist_comp", "dictorsetmaker", "exprlist", "testlist", "testlist_star_expr", "subscriptlist"}
)
_BRACKETS = frozenset("([{")
_CLOSERS = frozenset(")]}")
def _trailing_group(node: NodeOrLeaf) -> BaseNode | None:
"""The bracketed group ``node``'s code ends with, or ``None`` for none.
The one a formatter breaks a long line at: the last bracket to open at the
top level and close on the final character. For ``hero = Sprite2D(...)``
that is the construction, and for ``self.add_child(Sprite2D(...))`` it is
the ``add_child`` call rather than the construction inside it -- which is
the outer bracket, and so the first one to break.
"""
while is_node(node):
children = _children(node)
if not children:
return None
last = children[-1]
if is_leaf(last) and last.type == "operator" and last.value in _CLOSERS:
return node
node = last
return None
def _bracket_groups(node: NodeOrLeaf) -> Iterator[BaseNode]:
"""Every bracketed group under ``node``, outermost first."""
if not is_node(node):
return
children = _children(node)
if children:
first, last = children[0], children[-1]
if is_leaf(first) and first.value in _BRACKETS and is_leaf(last) and last.value in _CLOSERS:
yield node
for child in children:
yield from _bracket_groups(child)
def _group_parts(group: BaseNode) -> tuple[BaseNode, list[list[NodeOrLeaf]], Leaf | None]:
"""``(container, items, trailing comma)`` for a bracketed group.
``container`` is the node whose ``children`` the items live in: the
``arglist``-like node for a group holding several, and the group itself for
a group holding one, which parso puts under the bracket directly.
An item is the *run* of nodes between two commas, not a single node, because
what parso gives between them is not always one: a dict entry arrives as its
key, its colon and its value side by side, and laying out a dict as though
each of those were an item of its own would break the line after the key.
"""
inner = _children(group)[1:-1]
if len(inner) == 1 and inner[0].type in _SEQUENCES and is_node(inner[0]):
container: BaseNode = inner[0]
elements = _children(container)
else:
container = group
elements = list(inner)
items: list[list[NodeOrLeaf]] = []
run: list[NodeOrLeaf] = []
for element in elements:
if is_leaf(element) and element.type == "operator" and element.value == ",":
if run:
items.append(run)
run = []
continue
run.append(element)
if run:
items.append(run)
last = elements[-1] if elements else None
trailing = last if last is not None and is_leaf(last) and last.value == "," else None
return container, items, trailing
def _can_relayout(stmt: BaseNode, group: BaseNode) -> bool:
"""Is this statement one whose whitespace the layout may rewrite?
Only the whitespace the layout owns -- between the brackets, before the
commas, before the closing brackets -- may carry a line break, and nothing
in the statement may carry a comment: a comment sits on a line, and moving
the line out from under it would move the comment away from what it
describes. A value spanning lines of its own is the other refusal, since
its width cannot be measured and its shape is not this layer's to choose.
A line a semicolon has joined several statements onto is a third: the group
being laid out belongs to one of them, and the width of the line is the
width of all of them, so the measurement would be of something other than
the statement the edit was made in.
"""
if len(line_statements(stmt)) > 1:
return False
owned: set[int] = set()
for nested in _bracket_groups(stmt):
container, items, _trailing = _group_parts(nested)
owned.add(id(_children(nested)[-1]))
for child in _children(container):
owned.add(id(child if is_leaf(child) else child.get_first_leaf()))
for item in items:
owned.add(id(_first_leaf(item[0])))
first_leaf = stmt.get_first_leaf()
for leaf in _iter_leaves(stmt):
if "#" in leaf.prefix and leaf is not first_leaf:
# The first leaf's prefix is the indent, and whatever comment sat
# on the line above: neither is inside the statement.
return False
if leaf.type == "newline":
# The statement's own terminator, which no layout moves.
continue
if "\n" in leaf.value:
return False
if "\n" in leaf.prefix and leaf is not first_leaf and id(leaf) not in owned:
return False
return True
def _flatten(group: BaseNode) -> None:
"""Put ``group`` and everything nested in it back on one line.
The layout is decided from the width of the flat form, so the flat form is
what the tree is put into first. Every later decision only adds line breaks
back, which makes laying a statement out twice give what laying it out once
gave.
Flat also means evenly spaced: nothing after an opening bracket or before a
closing one, nothing before a comma and one space after it. That is the
spacing the width is measured at, so it has to be the spacing that is
written, and it is the spacing a formatter would have left anyway.
"""
for nested in _bracket_groups(group):
container, items, _trailing = _group_parts(nested)
_set_prefix(_children(nested)[-1], "")
for index, item in enumerate(items):
_set_prefix(item[0], "" if index == 0 else " ")
for child in _children(container):
if is_leaf(child) and child.type == "operator" and child.value == ",":
child.prefix = ""
def _first_leaf(node: NodeOrLeaf) -> Leaf:
return node if is_leaf(node) else node.get_first_leaf()
def _measure(leaves: list[Leaf], group: BaseNode, indent: int) -> tuple[int, int]:
"""Columns before ``group``'s opening bracket and after its closing one.
Both share the line the group starts on, and both count against the limit:
a construction is wrapped for the room left by the ``hero = `` in front of
it and the ``)`` of the call it may sit inside. ``leaves`` is what stands on
that line, in order: the statement's own, or one item's when the group is
nested inside it.
"""
open_leaf = _children(group)[0]
close_leaf = _children(group)[-1]
first_leaf = leaves[0] if leaves else None
before: list[str] = []
after: list[str] = []
state = 0
for leaf in leaves:
if leaf.type == "newline":
continue
piece = ("" if leaf is first_leaf else leaf.prefix) + leaf.value
if state == 0:
before.append(piece)
if leaf is open_leaf:
state = 1
elif state == 1:
if leaf is close_leaf:
state = 2
else:
after.append(piece)
head = "".join(before).rsplit("\n", 1)[-1]
return indent + display_width(head), display_width("".join(after))
def _lay_out_group(group: BaseNode, *, indent: int, prefix_width: int, suffix_width: int, limit: int) -> None:
"""Break ``group`` across lines as :func:`layout.wrap_statement` would.
Recurses into an item that is still too long on the line it was given, so a
material nested inside a construction is broken the same way the
construction was.
"""
container, items, trailing = _group_parts(group)
if not items:
return
texts = ["".join(node.get_code() for node in item).strip() for item in items]
collection = group.type != "trailer"
opener = _children(group)[0]
syntax = trailing_comma_is_syntax(
opener.value if is_leaf(opener) else "(", collection=collection, item_count=len(items)
)
explode = (trailing is not None and not syntax) or any(carries_magic_comma(text) for text in texts)
layout = call_layout(
texts,
indent=indent,
prefix_width=prefix_width,
suffix_width=suffix_width,
limit=limit,
explode=explode,
collection=collection,
extra=1 if trailing is not None else 0,
)
if layout == "line":
return
inner = "\n" + " " * (indent + INDENT)
for index, item in enumerate(items):
_set_prefix(item[0], " " if layout == "row" and index else inner)
_set_prefix(_children(group)[-1], "\n" + " " * indent)
if layout == "row":
return
if len(items) > 1 and trailing is None:
comma = _make_op(",", prefix="")
comma.parent = container
_children(container).append(comma)
# Each item now owns a line of its own, ending in the comma that separates
# it from the next, so what is left of the limit for it is one column less.
item_suffix = 1 if len(items) > 1 or trailing is not None else 0
for item in items:
nested = _trailing_group(item[-1])
if nested is None:
continue
leaves = [leaf for node in item for leaf in _iter_leaves(node)]
nested_prefix, _ = _measure(leaves, nested, indent + INDENT)
_lay_out_group(
nested, indent=indent + INDENT, prefix_width=nested_prefix, suffix_width=item_suffix, limit=limit
)
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _iter_leaves(node: NodeOrLeaf) -> Iterator[Leaf]:
"""Every leaf under ``node``, in source order."""
if is_leaf(node):
yield node
elif is_node(node):
for child in _children(node):
yield from _iter_leaves(child)
def _require_parent(node: NodeOrLeaf, op: str) -> BaseNode:
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: BaseNode, old: NodeOrLeaf, new: NodeOrLeaf) -> None:
children = _children(parent)
idx = _index_of(children, old)
new.parent = parent
old.parent = None
children[idx] = new
def _get_prefix(node: NodeOrLeaf) -> str:
if is_leaf(node):
return node.prefix
leaf: Leaf = node.get_first_leaf()
return 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) -> BaseNode | 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_node(node) and _is_call_trailer(node):
return node
if node.type == "atom_expr" and is_node(node):
last = _children(node)[-1]
if last.type == "trailer" and is_node(last) and _is_call_trailer(last):
return last
return None
def _is_call_trailer(trailer: BaseNode) -> bool:
if not trailer.children:
return False
first = _children(trailer)[0]
return is_leaf(first) and first.type == "operator" and first.value == "("
def _iter_arguments(trailer: BaseNode) -> Iterator[BaseNode]:
"""Yield ``argument`` nodes inside a call trailer."""
inner = _children(trailer)[1:-1] # strip ( and )
if not inner:
return
if len(inner) == 1:
node = inner[0]
if node.type == "argument" and is_node(node):
yield node
elif node.type == "arglist" and is_node(node):
for child in _children(node):
if child.type == "argument" and is_node(child):
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" and is_node(child):
yield child
def _argument_name(arg: BaseNode) -> str | None:
"""Return the kwarg name of an ``argument`` node, or None for positionals."""
children = _children(arg)
if len(children) < 3:
return None
name_node, eq, _value = children[0], children[1], children[2]
if name_node.type != "name" or eq.type != "operator" or not is_leaf(eq) or eq.value != "=":
return None
return name_node.value if is_leaf(name_node) else None
def _argument_value(arg: BaseNode) -> NodeOrLeaf | None:
children = _children(arg)
if len(children) < 3:
return None
return children[2]
def _argument_separator(trailer: BaseNode) -> str:
"""The prefix a newly appended argument should carry.
A space, for a call written on one line. For a call already written one
argument per line, the same line break and indent its arguments have, so
the new one lands where a reader expects it rather than trailing off the
end of the last one. Only the indent is copied: whatever comment sits in
that prefix belongs to the argument it was written above.
"""
last: NodeOrLeaf | None = None
for arg in _iter_arguments(trailer):
last = arg
if last is None:
return " "
prefix = _get_prefix(last)
return "\n" + _indent_of(prefix) if "\n" in prefix else " "
def _append_kwarg(trailer: BaseNode, name: str, value: NodeOrLeaf) -> None:
"""Append ``name=value`` to a call's arglist, preserving trailing comma."""
children = _children(trailer)
open_paren = children[0]
close_paren = children[-1]
inner = children[1:-1]
separator = _argument_separator(trailer)
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 != "arglist":
# Single-argument call: promote to arglist. Whatever the one argument
# is -- a named one, or a bare expression passed by position, which
# parso leaves sitting under the bracket as itself.
existing = inner[0]
comma = _make_op(",", prefix="")
_set_prefix(new_arg, separator)
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" and is_node(inner[0]):
arglist = inner[0]
arglist_children = _children(arglist)
last = arglist_children[-1]
had_trailing_comma = last.type == "operator" and is_leaf(last) and last.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)
comma = _make_op(",", prefix="")
comma.parent = arglist
arglist_children.insert(-1, comma)
arglist_children.insert(-1, new_arg)
else:
comma = _make_op(",", prefix="")
_set_prefix(new_arg, separator)
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) -> BaseNode:
"""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.
if not is_node(snippet):
raise ValueError("_build_kwarg_argument: snippet is not a composite node")
trailer = _children(snippet)[-1]
if not is_node(trailer):
raise ValueError("_build_kwarg_argument: trailer is not a composite node")
inner = _children(trailer)[1:-1]
if len(inner) == 1 and inner[0].type == "argument" and is_node(inner[0]):
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.
arg_children = _children(arg)
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: str = 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 is_leaf(node):
if node.type == "operator" and node.value == value:
return node
return None
if is_node(node):
for c in _children(node):
found = _find_op_leaf(c, value)
if found is not None:
return found
return None
def _make_arglist(args: list[NodeOrLeaf]) -> BaseNode:
"""Construct an ``arglist`` parso node from a list of ``argument``+``,``."""
snippet = parse_snippet("f(a, b)")
if not is_node(snippet):
raise ValueError("_make_arglist: snippet is not a composite node")
trailer = _children(snippet)[-1]
if not is_node(trailer):
raise ValueError("_make_arglist: trailer is not a composite node")
arglist = _children(trailer)[1]
if arglist.type != "arglist" or not is_node(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