Move root lattice definition to lattices dict (#51)

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

@@ -6,14 +6,17 @@
 
 from . import __version__, io, parse
 from .format import format_json
+from .migrate import MAX_VERSION
 from .migrate import migrate as _migrate
-from .validate import validate_file
+from .validate import parse_version, schema, validate_file
 
 FORMATS = "json", "lte", "madx"
 
 
 @click.group(context_settings=dict(max_content_width=120))
-@click.version_option(__version__)
+@click.version_option(
+    message=(f"LatticeJSON CLI, version {__version__}\n{schema['title']}")
+)
 def cli():
     pass
 
@@ -71,16 +74,30 @@ 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()

latticejson/convert.py

@@ -5,6 +5,7 @@
 
 from .exceptions import UnknownAttributeWarning, UnknownElementTypeWarning
 from .parse import parse_elegant, parse_madx
+from .validate import schema_version
 
 NAME_MAP = json.loads((Path(__file__).parent / "map.json").read_text())["map"]
 TO_ELEGANT = {x: y[0][0] for x, *y in NAME_MAP}
@@ -58,12 +59,14 @@ def _map_names(lattice_data: dict, name_map: dict):
                 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,
     )
 
 
@@ -74,21 +77,20 @@ 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
     for name, (type_, attributes) in elements.items():
         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 in sort_lattices(lattices):
+        strings.append(lattice_template(name, ", ".join(lattices[name])))
 
-    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
-    strings.append("USE, __MAIN__\n")
+    strings.append(f"USE, {latticejson['root']}\n")
     return "\n".join(strings)
 
 
@@ -99,21 +101,20 @@ 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
     for name, (type_, attributes) in elements.items():
         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 in sort_lattices(lattices):
+        strings.append(lattice_template(name, ", ".join(lattices[name])))
 
-    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
-    strings.append("USE, __MAIN__;\n")
+    strings.append(f"USE, {latticejson['root']};\n")
     return "\n".join(strings)
 
 

latticejson/exceptions.py

@@ -22,6 +22,13 @@ def __init__(self, object_name, lattice_name):
         )
 
 
+class IncompatibleVersionError(Exception):
+    """Raised if installed LatticeJSON library is out of date"""
+
+    def __init__(self, message):
+        super().__init__(f"Incompatible LatticeJSON version: {message}")
+
+
 class UnknownElementTypeWarning(UserWarning):
     """Raised if there is no equivalent LatticeJSON element."""
 

latticejson/migrate.py

@@ -1,24 +1,33 @@
-import sys
+def migrate(data: dict, initial: int, final: int) -> dict:
+    data = data.copy()
+    for from_, to, func in _VERSION_MAPS:
+        if from_ >= initial and (final is None or to <= final):
+            func(data)
+    return data
 
 
-def migrate(initial: dict, initial_version: tuple, final_version: tuple):
-    function = getattr(
-        sys.modules[__name__],
-        f"migrate_{'_'.join(initial_version)}_to_{'_'.join(final_version)}",
-        None,
-    )
-    if function is None:
-        raise NotImplementedError(f"Unkown versions {initial_version}, {final_version}")
+def _0_to_1(data: dict):
+    data["version"] = "1.0"
+    elements = data["elements"]
+    for name, attributes in elements.items():
+        elements[name] = attributes.pop("type"), attributes
 
-    return function(initial)
 
+def _1_to_2(data: dict):
+    data["version"] = "2.0"
+    data["title"] = data.pop("name")
+    data["lattices"] = data.pop("sub_lattices")
+    data["lattices"]["__MAIN__"] = data.pop("lattice")
+    data["root"] = "__MAIN__"
+    info = data.pop("description", False)
+    if info:
+        data["info"] = info
 
-def migrate_0_0_2_to_0_0_3(initial: dict):
-    final = initial.copy()
-    elements_final = {}
-    for name, attributes in final["elements"].items():
-        type_ = attributes.pop("type")
-        elements_final[name] = type_, attributes
+    for _, attributes in data["elements"].values():
+        info = attributes.pop("description", False)
+        if info:
+            attributes["info"] = info
 
-    final["elements"] = elements_final
-    return final
+
+_VERSION_MAPS = (0, 1, _0_to_1), (1, 2, _1_to_2)
+MAX_VERSION = _VERSION_MAPS[-1][1]