Python API¶

def format_statements(
    statements: Sequence[Statement] | Statement,
    options: FormatOptions = default_options,
) -> str:
    """Format a statement and return the code string"""
    if isinstance(statements, Statement):
        statements = [statements]

    res: list[OutputLine] = []

    def maybe_add_blank_line():
        if res and not res[-1].parts:
            return
        res.append(OutputLine())

    last_statement = None
    for statement in statements:
        if options.newline_before_new_type:
            if last_statement is not None:
                if (
                    options.newline_between_lines
                    and isinstance(statement, Line)
                    and isinstance(last_statement, Line)
                ):
                    maybe_add_blank_line()

                elif not isinstance(statement, type(last_statement)):
                    maybe_add_blank_line()
                elif (
                    isinstance(statement, Simple)
                    and statement.statement != last_statement.statement
                ):
                    maybe_add_blank_line()

        for line in format_nodes(statement, options=options):
            if not line.parts and not line.comment:
                maybe_add_blank_line()
            else:
                res.append(line)

        last_statement = statement

    while res and not res[0].parts and not res[0].comment:
        res = res[1:]

    text = "\n".join(line.render(options) for line in res)
    if options.newline_at_eof and text:
        return text + "\n"
    return text

def build_files(
    filenames: Sequence[str | pathlib.Path],
    *,
    combine: bool = False,
    root_path: pathlib.Path | None = None,
    input_format: str | None = None,
) -> list[Files]:
    """
    Construct one or more `Files` objects from CLI-style filename arguments.

    Parameters
    ----------
    filenames : list of str or Path
        Filenames to load. ``"-"`` reads from stdin.
    combine : bool, optional
        If True, all filenames are combined into a single `Files`
        (or `MemoryFiles` if any entry is stdin). If False (default),
        each filename becomes its own `Files`, preserving the per-file
        semantics of the legacy CLI loop.
    root_path : pathlib.Path, optional
        Directory used to resolve the synthetic stdin path. Defaults to ``Path.cwd()``.
    input_format : {"bmad", "namelist"}, optional
        Force how inputs are interpreted. ``None`` (default) auto-detects: an
        input is treated as a Tao ``tao.init`` namelist when it is named
        ``*.init`` or its contents look like a namelist (see
        `looks_like_namelist`); otherwise it is treated as Bmad. ``"namelist"``
        forces namelist handling for files that would not otherwise be detected;
        ``"bmad"`` forces lattice handling even for a ``*.init`` file.

    Returns
    -------
    list of Files
        One element if ``combine`` is True, otherwise one per input filename.
    """

    if not filenames:
        return []
    if root_path is None:
        root_path = pathlib.Path.cwd()

    def _is_stdin(fn) -> bool:
        return str(fn) == STDIN_TOKEN

    def _is_namelist(fn, contents: str | None = None) -> bool:
        """Decide whether ``fn`` should be handled as a Tao namelist file."""
        if input_format is not None:
            return input_format == "namelist"
        if is_init_file(fn):
            return True
        if contents is None:
            # Peek at the file to auto-detect a misnamed namelist.
            try:
                contents = pathlib.Path(fn).read_text()
            except OSError:
                return False
        return looks_like_namelist(contents)

    def _make_one(fn: str | pathlib.Path) -> Files:
        if _is_stdin(fn):
            fake_name = (root_path / STDIN_FAKE_NAME).resolve()
            contents = sys.stdin.read()
            if _is_namelist(fn, contents):
                files: Files = MemoryFiles.from_tao_init_contents(contents, fake_name)
            else:
                files = MemoryFiles(top_files=[fake_name], initial_contents={fake_name: contents})
            files.local_file_to_source_filename[fake_name] = STDIN_LABEL
            return files
        if _is_namelist(fn):
            return Files.from_tao_init(fn)
        return Files(top_files=[pathlib.Path(fn)])

    if not combine:
        return [_make_one(fn) for fn in filenames]

    # Combined mode: a single Files (or MemoryFiles if any stdin entry).
    stdin_path: pathlib.Path | None = None
    top_files: list[pathlib.Path] = []
    initial_contents: dict[pathlib.Path, str] = {}

    for fn in filenames:
        if _is_stdin(fn):
            if stdin_path is not None:
                raise ValueError("stdin ('-') can only be used once when combining inputs")
            stdin_path = (root_path / STDIN_FAKE_NAME).resolve()
            contents = sys.stdin.read()
            if _is_namelist(fn, contents):
                stdin_files = MemoryFiles.from_tao_init_contents(contents, stdin_path)
                top_files.extend(stdin_files.top_files)
                initial_contents.update(stdin_files.initial_contents)
            else:
                top_files.append(stdin_path)
                initial_contents[stdin_path] = contents
        elif _is_namelist(fn):
            top_files.extend(Files.from_tao_init(fn).top_files)
        else:
            top_files.append(pathlib.Path(fn))

    if initial_contents:
        files = MemoryFiles(top_files=top_files, initial_contents=initial_contents)
        if stdin_path is not None:
            files.local_file_to_source_filename[stdin_path] = STDIN_LABEL
        return [files]

    return [Files(top_files=top_files)]

def match_element_selector(statements: Iterable[Statement], selector: str) -> list[Element] | None:
    """
    Elements matched by a Bmad element selector, in definition order.

    Supported: ``*`` and ``%`` wildcards in the element name, and an optional
    ``class::pattern`` prefix whose class part (which may also use wildcards)
    is matched against each element's resolved type.  A plain name matches
    exactly.  All matching is case-insensitive.

    Returns None for selector syntax that is not supported yet:
    ranges (``q1:q5``), branch qualifiers (``lat>>q1``), instance counts
    (``q1##2``), and s-position selectors.
    """
    # TODO: branch qualifiers ("lat>>q1") and instance counts ("q1##2")
    if ">>" in selector or "##" in selector:
        return None

    class_pattern: str | None
    match selector.split("::"):
        case [name_pattern]:
            class_pattern = None
        case [class_pattern, name_pattern]:
            pass
        case _:
            return None
    # TODO: ranges ("q1:q5") and s-position selectors
    if ":" in (class_pattern or "") or ":" in name_pattern:
        return None

    matched = []
    for st in statements:
        if not isinstance(st, Element):
            continue
        if class_pattern is not None and (
            st.element_type is None or not _selector_part_matches(class_pattern, st.element_type)
        ):
            continue
        if _selector_part_matches(name_pattern, str(st.name)):
            matched.append(st)
    return matched

def target_selector_text(target: Token | Seq) -> str:
    """Reconstruct the text of an attribute-set statement's target (case-normalized)."""
    if isinstance(target, Token):
        return str(target)
    return str(target.to_token())

def calculate_diff(files1: Files, files2: Files) -> LatticeDiff:
    """
    Compute differences between two file sets and return a dataclass.

    Parameters
    ----------
    files1 : Files
        Left-hand file set.
    files2 : Files
        Right-hand file set.

    Returns
    -------
    LatticeDiff
        The structured differences.
    """
    diff = LatticeDiff()

    params1 = _collect_parameters(files1)
    params2 = _collect_parameters(files2)

    p_keys1 = set(params1.keys())
    p_keys2 = set(params2.keys())

    # Added
    for key in sorted(p_keys2 - p_keys1):
        diff.params_added.append(
            ParameterChange(key[0], key[1], old_value=None, new_value=params2[key])
        )

    # Removed
    for key in sorted(p_keys1 - p_keys2):
        diff.params_removed.append(
            ParameterChange(key[0], key[1], old_value=params1[key], new_value=None)
        )

    # Changed
    for key in sorted(p_keys1 & p_keys2):
        if params1[key] != params2[key]:
            diff.params_changed.append(
                ParameterChange(key[0], key[1], old_value=params1[key], new_value=params2[key])
            )

    elements1 = _collect_elements(files1)
    elements2 = _collect_elements(files2)

    e_keys1 = set(elements1.keys())
    e_keys2 = set(elements2.keys())

    diff.eles_added = sorted(e_keys2 - e_keys1)
    diff.eles_removed = sorted(e_keys1 - e_keys2)
    common_eles = e_keys1 & e_keys2

    def is_same_ele(ele1, ele2):
        e1 = elements1[ele1]
        e2 = elements2[ele2]
        return e1["type"] == e2["type"] and e1["attributes"] == e2["attributes"]

    diff.eles_renamed = [
        (ele1, ele2)
        for ele1 in diff.eles_removed
        for ele2 in diff.eles_added
        if is_same_ele(ele1, ele2)
    ]
    # Could detect multiple renames; A -> B1, B2
    # Not technically valid as far as a rename goes;
    # Make this instead a remove/add? Hmm
    for ele1, ele2 in diff.eles_renamed:
        try:
            diff.eles_removed.remove(ele1)
        except ValueError:
            pass
        try:
            diff.eles_added.remove(ele2)
        except ValueError:
            pass

    for name in common_eles:
        e1 = elements1[name]
        e2 = elements2[name]

        details = ElementDiffDetails()

        if e1["type"] != e2["type"]:
            details = dataclasses.replace(details, type_change=(e1["type"], e2["type"]))

        attrs1 = e1["attributes"]
        attrs2 = e2["attributes"]

        a_keys1 = set(attrs1.keys())
        a_keys2 = set(attrs2.keys())

        for a in sorted(a_keys2 - a_keys1):
            details.attrs_added.append((a, attrs2[a]))

        for a in sorted(a_keys1 - a_keys2):
            details.attrs_removed.append((a, attrs1[a]))

        for a in sorted(a_keys1 & a_keys2):
            if attrs1[a] != attrs2[a]:
                details.attrs_changed.append((a, attrs1[a], attrs2[a]))

        if details.has_changes:
            diff.eles_changed[name] = details

    return diff

def print_diff(diff: LatticeDiff, console: Console) -> None:
    """
    Render the diff using Rich tables.
    """
    if diff.has_param_diffs:
        console.rule("[bold]Parameters[/bold]")
        table = Table(show_header=True, header_style="bold magenta")
        table.add_column("State")
        table.add_column("Target")
        table.add_column("Name")
        table.add_column("Value (Left)", style="red")
        table.add_column("Value (Right)", style="green")

        for p in diff.params_added:
            table.add_row("Added", p.target, p.name, "", p.value_new_str, style="green")

        for p in diff.params_removed:
            table.add_row("Removed", p.target, p.name, p.value_old_str, "", style="red")

        for p in diff.params_changed:
            table.add_row(
                "Changed", p.target, p.name, p.value_old_str, p.value_new_str, style="yellow"
            )

        console.print(table)
        console.print()

    if diff.has_ele_diffs:
        console.rule("[bold]Elements[/bold]")

        table = Table(show_header=True, header_style="bold cyan")
        table.add_column("State")
        table.add_column("Element")
        table.add_column("Property/Attribute")
        table.add_column("Value (Left)", style="red")
        table.add_column("Value (Right)", style="green")

        for name in diff.eles_added:
            table.add_row("Added", name, "Element", "", "Exist", style="green")

        for name in diff.eles_removed:
            table.add_row("Removed", name, "Element", "Exist", "", style="red")

        for from_, to in diff.eles_renamed:
            table.add_row("Renamed", from_, "Element", from_, to, style="red")

        for name in sorted(diff.eles_changed.keys()):
            details = diff.eles_changed[name]

            if details.type_change:
                old_t, new_t = details.type_change
                table.add_row("Changed", name, "Type", old_t, new_t, style="magenta")

            for attr, val in details.attrs_added:
                table.add_row("Changed", name, f"Attr: {attr}", "", val, style="green")

            for attr, val in details.attrs_removed:
                table.add_row("Changed", name, f"Attr: {attr}", val, "", style="red")

            for attr, old_v, new_v in details.attrs_changed:
                table.add_row("Changed", name, f"Attr: {attr}", old_v, new_v, style="yellow")

        console.print(table)

def annotate_controller_variables(element: Element) -> None:
    """
    Annotate a controller's ``var={...}`` names and their usages.

    The declared variables are in ``var={...}``.

    They are then annotated in the control expressions, and as default-value
    attributes.
    """
    from .walk import iter_tokens

    var_names = {var._upper for var in get_controller_variables(element)}
    if not var_names:
        return

    var_attr = element.get_named_attribute("var", partial_match=False)

    def get_tokens():
        yield from iter_tokens(element.ele_list)
        for attr in element.attributes:
            yield from iter_tokens(attr)

    for tok in get_tokens():
        if (tok.role is None or tok.role == Role.name_) and tok._upper in var_names:
            tok.role = Role.controller_variable

    for attr in element.attributes:
        if attr is var_attr:
            continue
        if isinstance(attr.name, Token) and attr.name._upper in var_names:
            attr.name.role = Role.controller_variable

def get_controller_variables(element: Element) -> set[Token]:
    """
    Get variable names declared in an element's ``var={...}``.
    """
    from .walk import iter_tokens

    try:
        var_attr = element.get_named_attribute("var", partial_match=False)
    except KeyError:
        return set()

    if not isinstance(var_attr.value, Seq):
        return set()

    return {tok for tok in iter_tokens(var_attr.value)}