ast_utils.py

from __future__ import annotations

import ast
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from textwrap import indent
from tokenize import COMMENT as COMMENT_TOKEN
from tokenize import generate_tokens
from typing import Any, Iterator

import click

from idom import html


def rewrite_changed_nodes(
    file: Path,
    source: str,
    tree: ast.AST,
    changed: list[ChangedNode],
) -> str:
    ast.fix_missing_locations(tree)

    lines = source.split("\n")

    # find closest parent nodes that should be re-written
    nodes_to_unparse: list[ast.AST] = []
    for change in changed:
        node_lineage = [change.node, *change.parents]
        for i in range(len(node_lineage) - 1):
            current_node, next_node = node_lineage[i : i + 2]
            if (
                not hasattr(next_node, "lineno")
                or next_node.lineno < change.node.lineno
                or isinstance(next_node, (ast.ClassDef, ast.FunctionDef))
            ):
                nodes_to_unparse.append(current_node)
                break
        else:  # pragma: no cover
            raise RuntimeError("Failed to change code")

    # check if an nodes to rewrite contain eachother, pick outermost nodes
    current_outermost_node, *sorted_nodes_to_unparse = list(
        sorted(nodes_to_unparse, key=lambda n: n.lineno)
    )
    outermost_nodes_to_unparse = [current_outermost_node]
    for node in sorted_nodes_to_unparse:
        if (
            not current_outermost_node.end_lineno
            or node.lineno > current_outermost_node.end_lineno
        ):
            current_outermost_node = node
            outermost_nodes_to_unparse.append(node)

    moved_comment_lines_from_end: list[int] = []
    # now actually rewrite these nodes (in reverse to avoid changes earlier in file)
    for node in reversed(outermost_nodes_to_unparse):
        # make a best effort to preserve any comments that we're going to overwrite
        comments = _find_comments(lines[node.lineno - 1 : node.end_lineno])

        # there may be some content just before and after the content we're re-writing
        before_replacement = lines[node.lineno - 1][: node.col_offset].lstrip()

        after_replacement = (
            lines[node.end_lineno - 1][node.end_col_offset :].strip()
            if node.end_lineno is not None and node.end_col_offset is not None
            else ""
        )

        replacement = indent(
            before_replacement
            + "\n".join([*comments, ast.unparse(node)])
            + after_replacement,
            " " * (node.col_offset - len(before_replacement)),
        )

        lines[node.lineno - 1 : node.end_lineno or node.lineno] = [replacement]

        if comments:
            moved_comment_lines_from_end.append(len(lines) - node.lineno)

    for lineno_from_end in sorted(list(set(moved_comment_lines_from_end))):
        click.echo(f"Moved comments to {file}:{len(lines) - lineno_from_end}")

    return "\n".join(lines)


@dataclass
class ChangedNode:
    node: ast.AST
    parents: Sequence[ast.AST]


def find_element_constructor_usages(
    tree: ast.AST, add_props: bool = False
) -> Iterator[ElementConstructorInfo]:
    changed: list[Sequence[ast.AST]] = []
    for parents, node in _walk_with_parent(tree):
        if not (isinstance(node, ast.Call)):
            continue

        func = node.func
        if isinstance(func, ast.Attribute) and (
            (isinstance(func.value, ast.Name) and func.value.id == "html")
            or (isinstance(func.value, ast.Attribute) and func.value.attr == "html")
        ):
            name = func.attr
        elif isinstance(func, ast.Name):
            name = func.id
        else:
            continue

        maybe_attr_dict_node: Any | None = None

        if name == "vdom":
            if len(node.args) == 0:
                continue
            elif len(node.args) == 1:
                maybe_attr_dict_node = ast.Dict(keys=[], values=[])
                if add_props:
                    node.args.append(maybe_attr_dict_node)
                else:
                    continue
            elif isinstance(node.args[1], (ast.Constant, ast.JoinedStr)):
                maybe_attr_dict_node = ast.Dict(keys=[], values=[])
                if add_props:
                    node.args.insert(1, maybe_attr_dict_node)
                else:
                    continue
            elif len(node.args) >= 2:
                maybe_attr_dict_node = node.args[1]
        elif hasattr(html, name):
            if len(node.args) == 0:
                maybe_attr_dict_node = ast.Dict(keys=[], values=[])
                if add_props:
                    node.args.append(maybe_attr_dict_node)
                else:
                    continue
            elif isinstance(node.args[0], (ast.Constant, ast.JoinedStr)):
                maybe_attr_dict_node = ast.Dict(keys=[], values=[])
                if add_props:
                    node.args.insert(0, maybe_attr_dict_node)
                else:
                    continue
            else:
                maybe_attr_dict_node = node.args[0]

        if not maybe_attr_dict_node:
            continue

        if isinstance(maybe_attr_dict_node, ast.Dict) or (
            isinstance(maybe_attr_dict_node, ast.Call)
            and isinstance(maybe_attr_dict_node.func, ast.Name)
            and maybe_attr_dict_node.func.id == "dict"
            and isinstance(maybe_attr_dict_node.func.ctx, ast.Load)
        ):
            yield ElementConstructorInfo(node, maybe_attr_dict_node, parents)

    return changed


@dataclass
class ElementConstructorInfo:
    call: ast.Call
    props: ast.Dict | ast.Call
    parents: Sequence[ast.AST]


def _find_comments(lines: list[str]) -> list[str]:
    iter_lines = iter(lines)
    return [
        token
        for token_type, token, _, _, _ in generate_tokens(lambda: next(iter_lines))
        if token_type == COMMENT_TOKEN
    ]


def _walk_with_parent(
    node: ast.AST, parents: tuple[ast.AST, ...] = ()
) -> Iterator[tuple[tuple[ast.AST, ...], ast.AST]]:
    parents = (node,) + parents
    for child in ast.iter_child_nodes(node):
        yield parents, child
        yield from _walk_with_parent(child, parents)