Python API¶

def load_sources(
    self,
    base: pathlib.Path | None = None,
    reader: Callable[[pathlib.Path], str | None] | None = None,
) -> None:
    """
    Resolve and load the auxiliary files named in ``&tao_start``.

    Parameters
    ----------
    base : pathlib.Path, optional
        Directory relative names resolve against. Defaults to the tao.init
        directory.
    reader : callable, optional
        ``path -> text | None`` hook used to read each resolved file
        (``None`` = "missing"). Defaults to reading from disk.
    """
    if base is None:
        base = self.filename.parent if self.filename is not None else pathlib.Path()
    if reader is None:
        reader = _read_if_exists

    for key in SOURCE_FILE_KEYS:
        name = self._start_value(key)
        if not name:
            continue
        path = pathlib.Path(os.path.expandvars(name.strip()))
        if not path.is_absolute():
            path = base / path
        path = path.resolve()
        text = reader(path)
        if text is not None:
            self.sources[key] = NamelistFile.parse(text, filename=path)

def namelists_for(self, namelist_name: str) -> list[Namelist]:
    """All ``namelist_name`` groups from their resolved source, in file order."""
    source = self._source_for(namelist_name)
    return source.namelists_by_name.get(namelist_name.lower(), [])

def to_token(self, include_opener: bool = True) -> Token:
    """
    Convert this Block to a single Token with merged location information.
    """
    return Token.join(self.flatten(include_opener=include_opener))

def to_call_name(self) -> CallName:
    """Convert Seq to a single Token."""
    match self.items:
        case Token() as name, Seq(opener="(") as args:
            return CallName(name=name, args=args)
    raise UnexpectedCallName(
        f"Expected function call pattern not matched: {self} at {self.loc}"
    )

def to_text(self, opts: FormatOptions | None = None) -> str:
    """Convert Seq to its full output representation."""
    from .output import FormatOptions, format_nodes

    if opts is None:
        opts = FormatOptions()
    lines = format_nodes([self], options=opts)
    return "\n".join(line.render(options=opts) for line in lines)

def to_token(self, include_opener: bool = True) -> Token:
    """Convert Seq to a single Token."""
    from .output import FormatOptions, format_nodes

    if not include_opener:
        nodes = Seq(items=self.items, delimiter=self.delimiter).to_output_nodes()
    else:
        nodes = self.to_output_nodes()

    def check_can_tokenize(seq: Seq):
        for item in seq.items:
            if isinstance(item, Attribute):
                raise ValueError("Unable to tokenize Attributes")
            elif isinstance(item, Seq):
                check_can_tokenize(item)

    check_can_tokenize(self)

    opts = FormatOptions()
    (line,) = format_nodes(list(nodes), options=opts)
    line.comment = None
    return Token(line.render(options=opts), loc=self.loc)

def annotate(self):
    """
    Resolve named items across all parsed files.
    """
    named = self.get_named_items()
    for fn in self.by_filename:
        self._annotate_file(fn, named)
    self._resolve_element_types_in_order()

def flatten_all(self, call: bool, inline: bool) -> dict[pathlib.Path, list[Statement]]:
    """Flatten each top-level file independently, keyed by its path."""
    return {top: self.flatten(call=call, inline=inline, top=top) for top in self.top_files}

@classmethod
def from_tao_init(cls, path: pathlib.Path | str) -> "Files":
    """Build a :class:`Files` over the lattices listed in a ``tao.init``.

    The ``design_lattice(i)%file`` entries of the ``&tao_design_lattice``
    group become the top-level files, resolved (with environment-variable
    expansion) relative to the ``tao.init`` file's directory.
    """
    path = pathlib.Path(path)
    tao_init = TaoInit.from_file(path)
    top_files = _resolve_lattice_paths(tao_init.lattice_files, path.parent)
    return cls(top_files=top_files, tao_init=tao_init)

def get_named_items(self) -> dict[Token, Statement]:
    """
    Aggregate named items from all files.
    """
    named_items = {}
    for statements in self.by_filename.values():
        new_items = get_named_items(statements)
        # TODO: potential for linting with redef
        named_items.update(new_items)

    if "BEGINNING" not in named_items:
        named_items["BEGINNING"] = Element(
            name=Token("BEGINNING", loc=implicit_location, role=Role.name_),
            keyword=Token(
                "BEGINNING_ELE",
                loc=implicit_location,
                role=Role.kind,
            ),
        )
    if "END" not in named_items:
        named_items["END"] = Element(
            name=Token("END", loc=implicit_location, role=Role.name_),
            keyword=Token("MARKER", loc=implicit_location, role=Role.kind),
        )

    if "PARAMETER" not in named_items:
        named_items["PARAMETER"] = Element(
            name=Token("PARAMETER", loc=implicit_location, role=Role.name_),
            keyword=Token("!PARAMETER", loc=implicit_location, role=Role.kind),
        )

    if "PARTICLE_START" not in named_items:
        named_items["PARTICLE_START"] = Element(
            name=Token("PARTICLE_START", loc=implicit_location, role=Role.name_),
            keyword=Token("!PARTICLE_START", loc=implicit_location, role=Role.kind),
        )

    if "PTC_COM" not in named_items:
        named_items["PTC_COM"] = Element(
            name=Token("PTC_COM", loc=implicit_location, role=Role.name_),
            keyword=Token("!PTC_COM", loc=implicit_location, role=Role.kind),
        )

    return named_items

def get_statements_in_order(
    self,
    *,
    repeat_called_files: bool = True,
    top_files: Sequence[pathlib.Path] | None = None,
) -> list[Statement]:
    """
    Get all statements in order, as evaluated by Bmad.

    ``top_files`` restricts the traversal to the call trees of the given
    top-level files (defaulting to all of them).
    """

    handled = set()
    active = set()

    def _flatten(fn):
        if fn in active:
            # Guard against call cycles
            return
        if fn in handled and not repeat_called_files:
            return
        handled.add(fn)
        active.add(fn)
        try:
            for st in self.by_filename.get(fn, []):
                yield st
                if is_call_statement(st):
                    called = st.metadata.get("local_path")
                    if called is not None:
                        yield from _flatten(called)
        finally:
            active.discard(fn)

    res = []
    for fn in self.top_files if top_files is None else top_files:
        res.extend(_flatten(fn))
    return res

def match_elements(self, pattern: str) -> list[Element] | None:
    """
    Element definitions across all loaded files matching an element selector.

    See `match_element_selector` for the supported syntax; returns None for
    selector syntax that is not supported yet.
    """
    return match_element_selector(
        (st for statements in self.by_filename.values() for st in statements),
        pattern,
    )

def parse(
    self,
    recurse: bool = True,
    raise_if_missing: bool = False,
    keep_blocks: bool = False,
):
    """
    Parse the top-level file(s) and optionally their dependencies recursively.

    Parameters
    ----------
    recurse : bool, optional
        Recurse into called lattice files.  Defaults to True.
    raise_if_missing : bool, optional
        For lattice files included by way of ``call`` statements,
        this flag will control whether `FileNotFoundError` is raised.
        If a top-level file is missing, `FileNotFoundError` is always raised.
    keep_blocks : bool, optional
        Store the intermediate `Block` objects in
        ``self.blocks_by_filename`` so callers (e.g. verbose debug output)
        don't have to re-tokenize.
    """
    if not self.top_files:
        raise ValueError("Files requires at least one top-level file in top_files")

    self.top_files = [p.resolve() for p in self.top_files]
    top_set = set(self.top_files)

    if not self.stack:
        # Seed the stack so the first top file is popped first.
        for top in reversed(self.top_files):
            self.stack.append((pathlib.Path(top.name), top.parent))
            self.local_file_to_source_filename.setdefault(top, top.name)

    # We need to track processed files to avoid infinite loops in circular refs
    processed = set(self.by_filename.keys())

    while self.stack:
        filename_part, parent_dir = self.stack.pop()

        # Resolve the full path based on the parent context
        # (Note: filename_part might already be absolute if it's the main entry from disk)
        if filename_part.is_absolute():
            full_path = filename_part
        else:
            full_path = parent_dir / filename_part

        # Optimization: skip if already parsed
        if full_path in processed:
            continue

        logger.debug("Parsing %s", full_path)
        processed.add(full_path)

        try:
            contents = self._get_file_contents(full_path)
        except FileNotFoundError:
            logger.error(
                f"Could not find file: {full_path} (parent={parent_dir} file={filename_part})"
            )
            # Top-level files must exist. Otherwise, missing included files
            # are optionally an error.
            if full_path in top_set or raise_if_missing:
                raise FileNotFoundError(
                    f"Could not find file: {full_path} (parent={parent_dir} file={filename_part})"
                ) from None
            continue

        if is_init_file(full_path) or looks_like_namelist(contents):
            logger.debug("Skipping non-lattice (namelist) file: %s", full_path)
            self.by_filename[full_path] = []
            continue

        try:
            if keep_blocks:
                blocks = tokenize(contents=contents, filename=full_path)
                self.blocks_by_filename[full_path] = blocks
                statements: list[Statement] = [b.parse() for b in blocks]
            else:
                # We don't annotate individually here, we do it in bulk later
                statements = self._parse_file(contents, full_path)
        except Exception as ex:
            if hasattr(ex, "add_note"):  # py 3.11+
                ex.add_note(f"Exception ocurred while parsing {full_path}")
            raise

        self.by_filename[full_path] = statements

        for st in statements:
            if is_call_statement(st):
                # assert isinstance(st, Simple)
                st.metadata["local_path"] = self._add_file_by_statement(
                    statement_filename=full_path, st=st
                )

        if not recurse:
            # Without recursion, still process remaining top-level files,
            # but drop anything pulled in via `call` from this file.
            self.stack = [item for item in self.stack if (item[1] / item[0]) in top_set]
            if not self.stack:
                break

    return self.by_filename

def reformat(self, options: FormatOptions) -> None:
    """
    Reformat all files in the collection.
    """
    from .output import format_statements

    if options.flatten_call:
        for top, statements in self.flatten_all(
            call=options.flatten_call, inline=options.flatten_inline
        ).items():
            formatted = format_statements(statements, options)
            self._write_reformatted(top, formatted)
        return

    for fn, statements in self.by_filename.items():
        formatted = format_statements(statements, options)
        self._write_reformatted(fn, formatted)

@classmethod
def from_contents(cls, contents: str, root_path: pathlib.Path | str) -> MemoryFiles:
    """
    Create a MemoryFiles instance from a single string.

    Parameters
    ----------
    contents : str
        The source code content.
    root_path : pathlib.Path | str
        A "virtual" path where this file supposedly lives, used for resolving
        relative calls to other files.

    Returns
    -------
    MemoryFiles
        The initialized object (call .parse() on it next).
    """
    path = pathlib.Path(root_path).resolve()
    return cls(top_files=[path], initial_contents={path: contents})

@classmethod
def from_mapping(cls, contents: dict[pathlib.Path | str, str]) -> MemoryFiles:
    """
    Create a MemoryFiles instance from multiple in-memory files.

    Keys are treated as top-level files in iteration order.
    """
    resolved = {pathlib.Path(path).resolve(): cts for path, cts in contents.items()}
    return cls(top_files=list(resolved.keys()), initial_contents=resolved)

@classmethod
def from_tao_init_contents(
    cls,
    contents: str,
    root_path: pathlib.Path | str,
    lattice_contents: dict[pathlib.Path | str, str] | None = None,
) -> MemoryFiles:
    """
    Create a MemoryFiles from in-memory ``tao.init`` contents.

    Parameters
    ----------
    contents : str
        The ``tao.init`` file contents.
    root_path : pathlib.Path | str
        The "virtual" path of the ``tao.init`` file; lattice entries resolve
        relative to its parent directory.
    lattice_contents : dict, optional
        In-memory contents for referenced lattice files, keyed by path or
        name (resolved the same way as the ``design_lattice`` entries). Any
        lattice not provided here is read from disk when parsed.

    Returns
    -------
    MemoryFiles
        The initialized object (call .parse() on it next).
    """
    tao_path = pathlib.Path(root_path).resolve()
    base = tao_path.parent
    tao_init = TaoInit.parse(contents, filename=tao_path)
    top_files = _resolve_lattice_paths(tao_init.lattice_files, base)
    initial = {
        _resolve_lattice_paths([key], base)[0]: text
        for key, text in (lattice_contents or {}).items()
    }
    tao_init.load_sources(base=base, reader=initial.get)
    return cls(
        top_files=top_files,
        initial_contents=initial,
        tao_init=tao_init,
    )

def get_superposition_settings(
    self,
    deferred: dict[Token, dict[str, Token | Seq | None]] | None = None,
) -> SuperpositionSettings:
    """
    Resolve an element's effective ``superimpose``/``ref`` settings.

    Inline attributes are read first; deferred ``element[attr] = value``
    statements (as collected by `get_deferred_element_attributes`) override
    them.
    """

    def as_logical(value: Token | Seq | None) -> bool:
        text = _attribute_value_text(value)
        return text is None or text.upper() not in FALSE_LOGICALS_UPPER

    enabled = False
    ref_value: Token | Seq | None = None
    have_ref = False

    for attr in self.attributes:
        if not isinstance(attr.name, Token):
            continue
        uname = attr.name._upper
        if uname == "SUPERIMPOSE":
            enabled = as_logical(attr.value)
        elif uname == "REF":
            ref_value = attr.value
            have_ref = True

    overrides = (deferred or {}).get(self.name.upper(), {})
    if "SUPERIMPOSE" in overrides:
        enabled = as_logical(overrides["SUPERIMPOSE"])
    if "REF" in overrides:
        ref_value = overrides["REF"]
        have_ref = True

    ref = _attribute_value_text(ref_value) if have_ref else None
    ref_token = ref_value if isinstance(ref_value, Token) else None
    return SuperpositionSettings(enabled=enabled, ref=ref, ref_token=ref_token)