Merge remote-tracking branch 'origin/master' into poetry

Author: felix-andreas

README.md

@@ -5,34 +5,36 @@ has a human readable syntax and is available in all common programming language.
 appropriate choice to characterize the magnetic lattice of a particle accelerator.
 
 ## Specification
+
 This repository contains the
 [Specification of LatticeJSON](https://github.com/andreasfelix/latticejson/blob/master/latticejson/schema.json)
 in form of a [JSON Schema](https://json-schema.org).
 
-
 ## Example
 
 A LatticeJSON file for a FODO lattice:
+
 ```json
 {
-  "name": "FODO_RING",
-  "description": "This is the simplest possible strong focusing lattice.",
+  "version": "2.0",
+  "title": "FODO Lattice",
+  "info": "This is the simplest possible strong focusing lattice.",
+  "root": "RING",
   "elements": {
     "D1": ["Drift", {"length": 0.55}],
     "Q1": ["Quadrupole", {"length": 0.2, "k1": 1.2}],
     "Q2": ["Quadrupole", {"length": 0.4, "k1": -1.2}],
     "B1": ["Dipole", {"length": 1.5, "angle": 0.392701, "e1": 0.1963505, "e2": 0.1963505}]
   },
-  "sub_lattices": {
-    "FODO": ["Q1", "D1", "B1", "D1", "Q2", "D1", "B1", "D1", "Q1"]
-  },
-  "lattice": ["FODO", "FODO", "FODO", "FODO", "FODO", "FODO", "FODO", "FODO"]
+  "lattices": {
+    "CELL": ["Q1", "D1", "B1", "D1", "Q2", "D1", "B1", "D1", "Q1"],
+    "RING": ["CELL", "CELL", "CELL", "CELL", "CELL", "CELL", "CELL", "CELL"]
+  }
 }
-
 ```
 
+## LatticeJSON CLI
 
-# LatticeJSON CLI
 [![Python Version](https://img.shields.io/pypi/pyversions/latticejson)](https://pypi.org/project/latticejson/)
 [![PyPI](https://img.shields.io/pypi/v/latticejson.svg)](https://pypi.org/project/latticejson/)
 [![CI](https://github.com/andreasfelix/latticejson/workflows/CI/badge.svg)](https://github.com/andreasfelix/latticejson/actions?query=workflow%3ACI)
@@ -41,34 +43,41 @@ This repository also contains a Python based command-line tool which is able val
 and convert LatticeJSON files into other common lattice file formats and vice versa.
 
 You can install and update it using pip or pipenv:
+
 ```sh
 pip install -U latticejson
 ```
 
 Validate a LatticeJSON file:
+
 ```sh
 latticejson validate /path/to/lattice.json
 ```
 
 Convert an elegant lattice file to the LatticeJSON format:
+
 ```sh
 latticejson convert --to json /path/to/lattice.lte
 ```
 
 Autoformat one or more LatticeJSON files:
+
 ```sh
 latticejson autoformat /path/to/lattice.json ...
 ```
 
 To activate Bash completion add
-```
+
+```sh
 eval "$(_LATTICEJSON_COMPLETE=source latticejson)"
 ```
 
 to your `.bashrc`. Or, create an activation script with:
-```
+
+```sh
 _LATTICEJSON_COMPLETE=source latticejson > latticejson-complete.sh
 ```
 
 ## License
+
 [GNU General Public License v3.0](https://github.com/andreasfelix/latticejson/blob/master/LICENSE)

latticejson/cli.py

@@ -1,21 +1,22 @@
-import click
+import itertools
 import json
 from pathlib import Path
-import itertools
 
-from . import __version__
-from .validate import validate_file
-from . import io
+import click
+
+from . import __version__, io, parse
 from .format import format_json
+from .migrate import MAX_VERSION
 from .migrate import migrate as _migrate
-from . import parse
-
+from .validate import parse_version, schema, validate_file
 
 FORMATS = "json", "lte", "madx"
 
 
-@click.group()
-@click.version_option(__version__)
+@click.group(context_settings=dict(max_content_width=120))
+@click.version_option(
+    message=(f"LatticeJSON CLI, version {__version__}\n{schema['title']}")
+)
 def cli():
     pass
 
@@ -35,9 +36,12 @@ def cli():
     type=click.Choice(FORMATS, case_sensitive=False),
     help="Destination format",
 )
-def convert(file, from_, to):
+@click.option(
+    "--validate/--no-validate", default=True, help="Whether to validate the input file."
+)
+def convert(file, from_, to, validate):
     """Convert FILE (path or url) to another lattice file format."""
-    click.echo(io.save_string(io.load(file, from_), to))
+    click.echo(io.save_string(io.load(file, from_, validate), to))
 
 
 @cli.command()
@@ -70,16 +74,68 @@ def autoformat(files, dry_run):
 
 
 @cli.command()
-@click.argument("file", type=click.Path(exists=True))
-@click.option("--from", "from_", required=True, help="Initial version")
-@click.option("--to", required=True, help="Final version")
-def migrate(file, from_, to):
+@click.argument("files", nargs=-1, type=click.Path(exists=True))
+@click.option(
+    "--final", type=int, default=MAX_VERSION, show_default=True, help="Final version."
+)
+@click.option(
+    "--dry-run",
+    "-d",
+    is_flag=True,
+    help="Don't write the files back, just output the formatted files.",
+)
+def migrate(files, final, dry_run):
     """Migrate old LatticeJSON files to newer versions."""
-    text = Path(file).read_text()
-    initial_version = from_.split(".")
-    final_version = to.split(".")
-    latticejson = _migrate(json.loads(text), initial_version, final_version)
-    click.echo(format_json(latticejson))
+    for path in itertools.chain.from_iterable(
+        path.rglob("*.json") if path.is_dir() else (path,) for path in map(Path, files)
+    ):
+        data = json.loads(path.read_text())
+        initial = parse_version(data["version"]).major
+        latticejson = _migrate(data, initial, final)
+        formatted = format_json(latticejson)
+        click.secho(f"Migrated {path} from version {initial} to {final}", bold=True)
+        if dry_run:
+            click.echo(formatted)
+        else:
+            path.write_text(formatted)
+
+
+@cli.group()
+def utils():
+    """Some useful utilities."""
+    pass
+
+
+@utils.command()
+@click.argument("file", type=click.Path(exists=True))
+@click.option("--lattice", "-l", type=str, help="Root lattice of tree.")
+@click.option(
+    "--format",
+    "format_",
+    type=click.Choice(FORMATS, case_sensitive=False),
+    help="Source format [optional, default: use file extension]",
+)
+def tree(file, lattice, format_):
+    """Print tree of elements for a given LatticeJSON file."""
+    from .utils import tree
+
+    data = io.load(file, format_, validate)
+    click.echo(tree(data, lattice))
+
+
+@utils.command()
+@click.argument("file", type=click.Path(exists=True))
+@click.option("--lattice", "-l", type=str, help="New root lattice [Default: current]")
+@click.option("--warn-unused", "-w", is_flag=True, help="Log removed lattices.")
+@click.option(
+    "--validate/--no-validate", default=True, help="Whether to validate the input file."
+)
+def remove_unused(file, lattice, warn_unused, validate):
+    """Remove unused objects from a LatticeJSON file."""
+    from .utils import remove_unused
+
+    data = remove_unused(io.load(file, validate=validate), lattice, warn_unused)
+    click.echo(format_json(data))
 
 
 @cli.group()

latticejson/convert.py

@@ -1,15 +1,18 @@
-from typing import List, Dict
-from pathlib import Path
 import json
+from pathlib import Path
+from typing import Dict, List
 from warnings import warn
+
+from .exceptions import UnknownAttributeWarning, UnknownElementTypeWarning
 from .parse import parse_elegant, parse_madx
-from .exceptions import UnknownElementWarning, UnknownAttributeWarning
+from .utils import sort_lattices
+from .validate import schema_version
 
 NAME_MAP = json.loads((Path(__file__).parent / "map.json").read_text())["map"]
-JSON_TO_ELE = {x: y[0][0] for x, *y in NAME_MAP}
-ELE_TO_JSON = {y: x for x, *tup in NAME_MAP for y in tup[0]}
-JSON_TO_MADX = {x: y[1][0] for x, *y in NAME_MAP}
-MADX_TO_JSON = {y: x for x, *tup in NAME_MAP for y in tup[1]}
+TO_ELEGANT = {x: y[0][0] for x, *y in NAME_MAP}
+FROM_ELEGANT = {y: x for x, *tup in NAME_MAP for y in tup[0]}
+TO_MADX = {x: y[1][0] for x, *y in NAME_MAP}
+FROM_MADX = {y: x for x, *tup in NAME_MAP for y in tup[1]}
 
 
 def from_elegant(string):
@@ -22,7 +25,7 @@ def from_elegant(string):
     :type str, optional
     :return: dict in LatticeJSON format
     """
-    return _map_names_from_elegant(parse_elegant(string))
+    return _map_names(parse_elegant(string), FROM_ELEGANT)
 
 
 def from_madx(string):
@@ -35,15 +38,7 @@ def from_madx(string):
     :type str, optional
     :return: dict in LatticeJSON format
     """
-    return _map_names_from_madx(parse_madx(string))
-
-
-def _map_names_from_madx(madx_dict: dict):
-    return _map_names(madx_dict, MADX_TO_JSON)
-
-
-def _map_names_from_elegant(elegant_dict: dict):
-    return _map_names(elegant_dict, ELE_TO_JSON)
+    return _map_names(parse_madx(string), FROM_MADX)
 
 
 def _map_names(lattice_data: dict, name_map: dict):
@@ -52,7 +47,7 @@ def _map_names(lattice_data: dict, name_map: dict):
         latticejson_type = name_map.get(other_type)
         if latticejson_type is None:
             elements[name] = ["Drift", {"length": other_attributes.get("L", 0)}]
-            warn(UnknownElementWarning(name, other_type), stacklevel=2)
+            warn(UnknownElementTypeWarning(name, other_type))
             continue
 
         attributes = {}
@@ -62,15 +57,17 @@ def _map_names(lattice_data: dict, name_map: dict):
             if latticejson_key is not None:
                 attributes[latticejson_key] = value
             else:
-                warn(UnknownAttributeWarning(other_key, name), stacklevel=2)
+                warn(UnknownAttributeWarning(other_key, name))
 
     lattices = lattice_data["lattices"]
-    lattice_name, main_lattice = lattices.popitem()  # use last lattice as main_lattice
+    root = lattice_data.get("root", tuple(lattices.keys())[-1])
+    title = lattice_data.get("title", "")
     return dict(
-        name=lattice_name,
-        lattice=main_lattice,
-        sub_lattices=lattices,
+        version=str(schema_version),
+        title=title,
+        root=root,
         elements=elements,
+        lattices=lattices,
     )
 
 
@@ -81,21 +78,21 @@ def to_elegant(latticejson: dict) -> str:
     :return: string with in elegant lattice file format
     """
     elements = latticejson["elements"]
-    sub_lattices = latticejson["sub_lattices"]
+    lattices = latticejson["lattices"]
 
-    strings = [f"! TITLE: {latticejson['name']}"]
+    strings = [f"! TITLE: {latticejson['title']}"]
     element_template = "{}: {}, {}".format
+    # TODO: check if equivalent type exists in elegant
     for name, (type_, attributes) in elements.items():
-        attrs = ", ".join(f"{JSON_TO_ELE[k]}={v}" for k, v in attributes.items())
-        elegant_type = JSON_TO_ELE[type_]
+        attrs = ", ".join(f"{TO_ELEGANT[k]}={v}" for k, v in attributes.items())
+        elegant_type = TO_ELEGANT[type_]
         strings.append(element_template(name, elegant_type, attrs))
 
     lattice_template = "{}: LINE=({})".format
-    for name in sort_lattices(sub_lattices):
-        strings.append(lattice_template(name, ", ".join(sub_lattices[name])))
+    for name, children in sort_lattices(latticejson).items():
+        strings.append(lattice_template(name, ", ".join(children)))
 
-    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
-    strings.append("USE, __MAIN__\n")
+    strings.append(f"USE, {latticejson['root']}\n")
     return "\n".join(strings)
 
 
@@ -106,39 +103,19 @@ def to_madx(latticejson: dict) -> str:
     :return: string with in elegant lattice file format
     """
     elements = latticejson["elements"]
-    sub_lattices = latticejson["sub_lattices"]
+    lattices = latticejson["lattices"]
 
-    strings = [f"TITLE, \"{latticejson['name']}\""]
+    strings = [f"TITLE, \"{latticejson['title']}\";"]
     element_template = "{}: {}, {};".format
+    # TODO: check if equivalent type exists in madx
     for name, (type_, attributes) in elements.items():
-        attrs = ", ".join(f"{JSON_TO_MADX[k]}={v}" for k, v in attributes.items())
-        elegant_type = JSON_TO_MADX[type_]
+        attrs = ", ".join(f"{TO_MADX[k]}={v}" for k, v in attributes.items())
+        elegant_type = TO_MADX[type_]
         strings.append(element_template(name, elegant_type, attrs))
 
     lattice_template = "{}: LINE=({});".format
-    for name in sort_lattices(sub_lattices):
-        strings.append(lattice_template(name, ", ".join(sub_lattices[name])))
+    for name, children in sort_lattices(latticejson).items():
+        strings.append(lattice_template(name, ", ".join(children)))
 
-    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
-    strings.append("USE, __MAIN__;\n")
+    strings.append(f"USE, SEQUENCE={latticejson['root']};\n")
     return "\n".join(strings)
-
-
-def sort_lattices(lattices: Dict[str, List[str]]) -> List[str]:
-    """Returns a sorted list of lattice names for a given dict of lattices."""
-
-    lattices_set = set(lattices)
-    lattice_names = []
-
-    def _sort_lattices(name):
-        for child_name in lattices[name]:
-            if child_name in lattices_set:
-                lattices_set.remove(child_name)
-                _sort_lattices(child_name)
-
-        lattice_names.append(name)
-
-    while len(lattices_set) > 0:
-        _sort_lattices(lattices_set.pop())
-
-    return lattice_names

latticejson/elegant.lark

@@ -10,7 +10,7 @@
 %ignore /!.*/            // ingore comments
 %ignore /[ \t\f]/+       // ingore whitespace
 %ignore /&[ \t\f]*\r?\n/ // line continuation
-%import common (SIGNED_INT, SIGNED_FLOAT, NUMBER, ESCAPED_STRING, CNAME)
+%import common (SIGNED_INT, SIGNED_FLOAT, SIGNED_NUMBER, ESCAPED_STRING, CNAME)
 
 int         : SIGNED_INT
 float       : SIGNED_FLOAT
@@ -38,7 +38,7 @@ command     : name ["," word]
 // identified as string.
 
 assignment  : expr "sto" CNAME
-?expr       : NUMBER        -> number
+?expr       : SIGNED_NUMBER -> number
             | CNAME         -> variable
             | function
             | binary