diff --git a/src/latform/dump.py b/src/latform/dump.py index 4cd1cb0..6651f52 100644 --- a/src/latform/dump.py +++ b/src/latform/dump.py @@ -222,7 +222,10 @@ def _resolve_used_elements( """ Resolve which named items are active in the expanded lattice. - Starting from the roots of the last USE statement, lines are expanded + Roots are resolved per top-level lattice file: the last USE statement in + each file's call tree, or the ``@use_line`` names from the corresponding + ``tao.init`` ``design_lattice(i)%file`` entry, which override that + lattice's own USE statement. From those roots, lines are expanded recursively (including repetitions, reflections, replacement-line calls, list members, and fork targets). A fixpoint pass then folds in usage that depends on other elements being used: @@ -239,17 +242,17 @@ def _resolve_used_elements( Mapping of uppercased active names to a human-readable reason. """ all_statements = files.get_statements_in_order(repeat_called_files=False) - use_cmds = [ - st - for st in all_statements - if isinstance(st, Simple) and st.statement._upper == "USE" and st.arguments - ] - # Bmad semantics: the last USE statement wins; each of its arguments is - # the root line of a branch. Bare argument tokens are parsed as - # value-less Attributes. - roots: list[Token] = [] - if use_cmds: + def last_use_roots(statements: list[Statement]) -> list[Token]: + # Bmad semantics: the last USE statement wins; each of its arguments is + # the root line of a branch. Bare argument tokens are parsed as + # value-less Attributes. + use_cmds = [ + st + for st in statements + if isinstance(st, Simple) and st.statement._upper == "USE" and st.arguments + ] + roots: list[Token] = [] for use_cmd in reversed(use_cmds): for arg in use_cmd.arguments: if isinstance(arg, Token): @@ -264,6 +267,27 @@ def _resolve_used_elements( # I think this is a scenario that's common in reused sublattices # that can be standalone break + return roots + + # Roots are resolved per top-level lattice file. A tao.init + # ``design_lattice(i)%file = 'lat.bmad@line_name'`` suffix overrides that + # lattice's own USE statement (Tao/bmad_parser semantics); lattices without + # a suffix fall back to their in-file USE. + tao_entries: list[tuple[str, list[str]]] = ( + files.tao_init.lattice_file_with_use_line if files.tao_init is not None else [] + ) + root_reasons: dict[Token, str] = {} + for index, top_file in enumerate(files.top_files): + use_lines = tao_entries[index][1] if index < len(tao_entries) else [] + if use_lines: + for name in use_lines: + root_reasons.setdefault(Token(name.upper()), "tao.init use_line") + else: + statements = files.get_statements_in_order( + repeat_called_files=False, top_files=[top_file] + ) + for root in last_use_roots(statements): + root_reasons.setdefault(root, "use statement") deferred = get_deferred_element_attributes(all_statements) used: dict[Token, str] = {} @@ -288,10 +312,10 @@ def mark(name: Token, reason: str) -> None: for target in _fork_targets(item, deferred): mark(target, f"fork target of {name}") - for root in roots: - mark(root, "use statement") + for root, reason in root_reasons.items(): + mark(root, reason) - if roots: + if root_reasons: # The implicit lattice endpoints exist in any expanded lattice for name in (Token("BEGINNING"), Token("END")): if name in named_items: @@ -319,7 +343,7 @@ def mark(name: Token, reason: str) -> None: changed = True continue - if not roots: + if not root_reasons: continue superposition = element.get_superposition_settings(deferred) diff --git a/src/latform/parser.py b/src/latform/parser.py index 9653458..48557ca 100644 --- a/src/latform/parser.py +++ b/src/latform/parser.py @@ -944,8 +944,18 @@ def flatten_all(self, call: bool, inline: bool) -> dict[pathlib.Path, list[State """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} - def get_statements_in_order(self, *, repeat_called_files: bool = True) -> list[Statement]: - """Get all statements in order, as evaluated by Bmad.""" + 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() @@ -969,7 +979,7 @@ def _flatten(fn): active.discard(fn) res = [] - for fn in self.top_files: + for fn in self.top_files if top_files is None else top_files: res.extend(_flatten(fn)) return res diff --git a/src/latform/tao/file.py b/src/latform/tao/file.py index c5d67ad..ed17262 100644 --- a/src/latform/tao/file.py +++ b/src/latform/tao/file.py @@ -421,21 +421,45 @@ def _lattice_file_assignments(self) -> list[Assignment]: ] @property - def lattice_files(self) -> list[str]: + def lattice_file_with_use_line(self) -> list[tuple[str, list[str]]]: """ - Ordered ``design_lattice(i)%file`` values (unquoted), by index. + Ordered, unquoted ``design_lattice(i)%file`` values with `use_line` names. - When set, rewritse the ``design_lattice(i)%file`` entries to ``files`` (1-based). - Existing entries are updated in place. Non-matching additional entries - are removed, and new entries are appended. + When set, rewrites the entries from ``(filename, use_line_names)`` + tuples, joining the names back into the ``@`` suffix form. """ + + def item(assignment: Assignment) -> tuple[str, list[str]]: + # Unquote and remove the '@use_line1@use_line2...' suffix + + filename = unquote_value(assignment.value.strip()) + if "@" in filename: + filename, use_lines = filename.split("@", 1) + return (filename, use_lines.split("@")) + return filename, [] + by_index = { - assignment.path.components[0].index: unquote_value(assignment.value.strip()) + assignment.path.components[0].index: item(assignment) for assignment in self._lattice_file_assignments() if assignment.path.components[0].index is not None } return [by_index[i] for i in sorted(by_index)] + @lattice_file_with_use_line.setter + def lattice_file_with_use_line(self, entries: list[tuple[str, list[str]]]) -> None: + self.lattice_files = ["@".join([fn, *use_lines]) for fn, use_lines in entries] + + @property + def lattice_files(self) -> list[str]: + """ + Ordered ``design_lattice(i)%file`` values (unquoted), by index. + + When set, rewrites the ``design_lattice(i)%file`` entries to ``files`` (1-based). + Existing entries are updated in place. Non-matching additional entries + are removed, and new entries are appended. + """ + return [fn for fn, _ in self.lattice_file_with_use_line] + @lattice_files.setter def lattice_files(self, files: list[str]) -> None: namelist = self.design_lattice diff --git a/src/latform/template.py b/src/latform/template.py index 320649c..cc91bd3 100644 --- a/src/latform/template.py +++ b/src/latform/template.py @@ -253,17 +253,6 @@ def _instantiate_tao_init( in_dir = posixpath.dirname(input_rel) out_dir = posixpath.dirname(output_rel) - remapped: list[str] = [] - changed = False - for entry in tao_init.lattice_files: - mapped = _map_reference(entry, in_dir, out_dir, explicit_paths, in_to_out) - if mapped is None: - remapped.append(entry) - else: - remapped.append(mapped) - changed = True - if changed: - tao_init.lattice_files = remapped merged: dict[str, str] = {} if rules and (rules["literal"] or rules["regex"] or rules["parts"]): @@ -276,6 +265,20 @@ def _instantiate_tao_init( if merged: rename_tao_elements(tao_init, merged) + remapped: list[tuple[str, list[str]]] = [] + changed = False + for filename, use_lines in tao_init.lattice_file_with_use_line: + mapped = _map_reference(filename, in_dir, out_dir, explicit_paths, in_to_out) + mapped = mapped or filename + + new_use_lines = [merged.get(name.upper(), name) for name in use_lines] + remapped.append((mapped, new_use_lines)) + + if mapped != filename or new_use_lines != use_lines: + changed = True + if changed: + tao_init.lattice_file_with_use_line = remapped + for name_key, assignments in ((override or {}).get("namelists") or {}).items(): name, index = split_namelist_key(name_key) interpolated = {k: _interpolate(str(v), instance) for k, v in assignments.items()} diff --git a/src/latform/tests/test_dump.py b/src/latform/tests/test_dump.py index 44eb007..4949a28 100644 --- a/src/latform/tests/test_dump.py +++ b/src/latform/tests/test_dump.py @@ -227,6 +227,69 @@ def test_reasons(): assert status["UNUSED_Q"]["reason"] == "" +TAO_TWO_LINE_LATTICE = """\ +m1: marker +m2: marker +l1: line = (m1) +l2: line = (m2) +use, l1 +""" + + +def _tao_files(entries: list[str], lattice_contents: dict[str, str]) -> MemoryFiles: + body = "\n".join( + f" design_lattice({i})%file = '{entry}'" for i, entry in enumerate(entries, start=1) + ) + files = MemoryFiles.from_tao_init_contents( + f"&tao_design_lattice\n{body}\n/\n", + "/virtual/tao.init", + lattice_contents=lattice_contents, + ) + files.parse() + files.annotate() + return files + + +def _tao_used(entries: list[str], lattice_contents: dict[str, str]) -> set[str]: + files = _tao_files(entries, lattice_contents) + return {row["name"] for row in get_elements_status(files, "all") if row["used"] == "YES"} + + +@pytest.mark.parametrize( + ("entry", "expected_used", "expected_unused"), + [ + pytest.param("mem.lat.bmad", {"L1", "M1"}, {"L2", "M2"}, id="no-use-line"), + pytest.param("mem.lat.bmad@l2", {"L2", "M2"}, {"L1", "M1"}, id="use-line-overrides-use"), + pytest.param( + "mem.lat.bmad@l1@l2", {"L1", "M1", "L2", "M2"}, set(), id="multiple-use-lines" + ), + ], +) +def test_tao_init_use_line_roots(entry: str, expected_used: set[str], expected_unused: set[str]): + used = _tao_used([entry], {"mem.lat.bmad": TAO_TWO_LINE_LATTICE}) + assert expected_used <= used + assert not expected_unused & used + + +def test_tao_init_use_line_is_per_file(): + """A use_line overrides only its own lattice's USE; other lattices keep theirs.""" + first = "m1: marker\nm2: marker\nl1: line = (m1)\nl2: line = (m2)\nuse, l1\n" + second = "m3: marker\nm4: marker\nl3: line = (m3)\nl4: line = (m4)\nuse, l3\n" + used = _tao_used( + ["first.lat.bmad", "second.lat.bmad@l4"], + {"first.lat.bmad": first, "second.lat.bmad": second}, + ) + assert {"L1", "M1", "L4", "M4"} <= used + assert not {"L2", "M2", "L3", "M3"} & used + + +def test_tao_init_use_line_reason(): + files = _tao_files(["mem.lat.bmad@l2"], {"mem.lat.bmad": TAO_TWO_LINE_LATTICE}) + status = {row["name"]: row for row in get_elements_status(files, "all")} + assert status["L2"]["reason"] == "tao.init use_line" + assert status["M2"]["reason"] == "in line L2" + + def test_get_constants(): src = """ my_const = 0.5 diff --git a/src/latform/tests/test_parse.py b/src/latform/tests/test_parse.py index d7edf34..3df2da3 100644 --- a/src/latform/tests/test_parse.py +++ b/src/latform/tests/test_parse.py @@ -2745,6 +2745,32 @@ def test_roundtrip(code: str, expected) -> None: roundtrip_code(code) +@pytest.mark.parametrize( + "code", + [ + pytest.param("q: quad, l = ((1+2))*3", id="doubled-parens"), + pytest.param("ov: overlay = {{q[k1]: 0.1}}, var = {a}", id="doubled-braces"), + pytest.param( + "foo[phi0_multipass] = -(acos(e2/e1*cos(phi/180*pi)) + phi/180*pi)/(2*pi)", + id="negated-group", + ), + pytest.param( + "bar[phi0_multipass] = -(\n acos(e2/e1*cos(phi/180*pi)) + phi/180*pi\n)/(2*pi)", + id="negated-group-multiline", + ), + ], +) +def test_nested_parens_roundtrip(code: str) -> None: + # Redundant nesting like (( )) is collapsed on output, so this checks + # parse -> format -> parse equality rather than byte preservation. + roundtrip_code(code) + + +def test_nested_brackets_invalid() -> None: + with pytest.raises(ValueError, match=r"Nested '\[\[ \]\]'"): + parse("m: marker, type = [[x]]") + + @pytest.mark.parametrize( ("code",), [ diff --git a/src/latform/tests/test_tao_init.py b/src/latform/tests/test_tao_init.py index 6566c31..639cf8c 100644 --- a/src/latform/tests/test_tao_init.py +++ b/src/latform/tests/test_tao_init.py @@ -88,6 +88,44 @@ def test_lattice_files_ordered_and_skips_comments(): ] +USE_LINE_INIT = """\ +&tao_design_lattice + design_lattice(2)%file = 'b.lat.bmad@lineA' + design_lattice(1)%file = 'a.lat.bmad' + design_lattice(3)%file = 'sub/c.lat.bmad@lineA@lineB' +/ +""" + + +def test_lattice_file_with_use_line(): + tao = TaoInit.parse(USE_LINE_INIT) + assert tao.lattice_file_with_use_line == [ + ("a.lat.bmad", []), + ("b.lat.bmad", ["lineA"]), + ("sub/c.lat.bmad", ["lineA", "lineB"]), + ] + # The '@use_line' suffixes are stripped from the plain filename view + assert tao.lattice_files == ["a.lat.bmad", "b.lat.bmad", "sub/c.lat.bmad"] + assert tao.render() == USE_LINE_INIT + + +def test_set_lattice_file_with_use_line(): + tao = TaoInit.parse(USE_LINE_INIT) + entries = [("a2.lat.bmad", []), ("b2.lat.bmad", ["lineC"]), ("c2.lat.bmad", ["l1", "l2"])] + tao.lattice_file_with_use_line = entries + assert tao.lattice_file_with_use_line == entries + assert tao.lattice_files == ["a2.lat.bmad", "b2.lat.bmad", "c2.lat.bmad"] + assert TaoInit.parse(tao.render()).lattice_file_with_use_line == entries + + +def test_set_lattice_files_passes_at_suffix_through(): + """The plain setter writes '@' suffixes verbatim; the getter strips them (asymmetric).""" + tao = TaoInit.parse(USE_LINE_INIT) + tao.lattice_files = ["x.lat.bmad@ln"] + assert tao.lattice_files == ["x.lat.bmad"] + assert tao.lattice_file_with_use_line == [("x.lat.bmad", ["ln"])] + + def test_keypath_decomposes_nested_key(): path = KeyPath.parse("foo(3)%bar(2)%val") assert path.names == ("foo", "bar", "val") @@ -250,6 +288,21 @@ def test_memory_files_from_tao_init_contents(): assert "M_Q" in files.get_named_items() +def test_memory_files_from_tao_init_use_line(): + """A '@use_line' suffix does not interfere with resolving the lattice file itself.""" + contents = "&tao_design_lattice\n design_lattice(1)%file = 'mem.lat.bmad@ml'\n/\n" + root = FILES / "virtual" / "tao.init" + files = MemoryFiles.from_tao_init_contents( + contents, + root, + lattice_contents={"mem.lat.bmad": "M_Q: quadrupole, l = 1\nml: line = (M_Q)\nuse, ml\n"}, + ) + assert [p.name for p in files.top_files] == ["mem.lat.bmad"] + files.parse() + files.annotate() + assert "M_Q" in files.get_named_items() + + @pytest.mark.parametrize( ("name", "expected"), [("tao.init", True), ("tao_plot.init", True), ("foo.INIT", True), ("lat.bmad", False)], diff --git a/src/latform/tests/test_template.py b/src/latform/tests/test_template.py index 489d375..3c35042 100644 --- a/src/latform/tests/test_template.py +++ b/src/latform/tests/test_template.py @@ -346,6 +346,24 @@ def test_instantiate_renames_elements_in_tao_init(tmp_path): assert "var(1)%ele_name='C1_Q'" in flat +def test_instantiate_renames_use_line_in_tao_init(tmp_path): + """A design_lattice '@use_line' suffix is kept and renamed along with the line.""" + (tmp_path / "cx.lat.bmad").write_text( + "CX_Q: quadrupole, k1=0.0\nCX_LINE: line=(CX_Q)\nuse, CX_LINE\n" + ) + (tmp_path / "tao.init").write_text( + "&tao_design_lattice\n design_lattice(1)%file = 'cx.lat.bmad@cx_line'\n/\n" + ) + spec = { + "template": [{"input": "cx.lat.bmad", "output": "{instance}.lat.bmad"}], + "renames": {r"CX(_.*|$)": r"{instance:upper}\1"}, + "tao_init": {"input": "tao.init", "output": "{instance}/tao.init"}, + "instances": {"c1": {}}, + } + out = instantiate(spec, base_dir=tmp_path, options=default_options)["c1"]["c1/tao.init"] + assert "design_lattice(1)%file = '../c1.lat.bmad@C1_LINE'" in out + + def test_instantiate_renames_tao_init_only_elements(tmp_path): """Names only the tao.init references (not defined in the template lattices) are still renamed by the instance's rename rules.""" diff --git a/src/latform/types.py b/src/latform/types.py index 6a5e78b..66ba67f 100644 --- a/src/latform/types.py +++ b/src/latform/types.py @@ -68,7 +68,12 @@ def _in_expected_delimiter( ) -> Self: unwrapped, inner_delimiter = self.maybe_unwrap() assert unwrapped is not None - assert not isinstance(unwrapped, Seq) + if isinstance(unwrapped, Seq): + if opener == "[" and unwrapped.opener == "[": + raise ValueError(f"Nested '[[ ]]' is not valid syntax at {opener.loc}") + # Nested sequences, possibly unnecessarily (TODO) + return unwrapped + if inner_delimiter == expected_delimiter or not self.items: # If the sequence is already delimited as expected, give it back directly items = unwrapped