Add some utils (#50)

* add print_tree
* new sort_lattices
* add removed unused function

Author: felix-andreas

latticejson/cli.py

@@ -100,6 +100,44 @@ def migrate(files, final, dry_run):
             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()
 def debug():
     """Some useful commands for debugging/development."""

latticejson/convert.py

@@ -5,6 +5,7 @@
 
 from .exceptions import UnknownAttributeWarning, UnknownElementTypeWarning
 from .parse import parse_elegant, parse_madx
+from .utils import sort_lattices
 from .validate import schema_version
 
 NAME_MAP = json.loads((Path(__file__).parent / "map.json").read_text())["map"]
@@ -87,8 +88,8 @@ def to_elegant(latticejson: dict) -> str:
         strings.append(element_template(name, elegant_type, attrs))
 
     lattice_template = "{}: LINE=({})".format
-    for name in sort_lattices(lattices):
-        strings.append(lattice_template(name, ", ".join(lattices[name])))
+    for name, children in sort_lattices(latticejson).items():
+        strings.append(lattice_template(name, ", ".join(children)))
 
     strings.append(f"USE, {latticejson['root']}\n")
     return "\n".join(strings)
@@ -111,28 +112,8 @@ def to_madx(latticejson: dict) -> str:
         strings.append(element_template(name, elegant_type, attrs))
 
     lattice_template = "{}: LINE=({});".format
-    for name in sort_lattices(lattices):
-        strings.append(lattice_template(name, ", ".join(lattices[name])))
+    for name, children in sort_lattices(latticejson).items():
+        strings.append(lattice_template(name, ", ".join(children)))
 
     strings.append(f"USE, {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/utils.py

@@ -0,0 +1,81 @@
+from itertools import chain
+from warnings import warn
+
+
+def tree(latticejson, name=None):
+    lattices = latticejson["lattices"]
+
+    def _tree(name, prefix=""):
+        string = f"{name}\n"
+        if name in lattices:
+            *other, last = lattices[name]
+            for child in other:
+                string += f"{prefix}├─── {_tree(child, prefix + '│   ')}"
+            string += f"{prefix}└─── {_tree(last, prefix + '    ')}"
+        return string
+
+    return _tree(latticejson["root"] if name is None else name)
+
+
+def sort_lattices(latticejson, root=None, keep_unused=False):
+    """Returns a sorted dict of lattice objects."""
+    lattices = latticejson["lattices"]
+    lattices_set = set(lattices)
+    lattices_sorted = {}
+
+    def _sort_lattices(name):
+        lattices_set.remove(name)
+        for child in lattices[name]:
+            if child in lattices_set:
+                _sort_lattices(child)
+        lattices_sorted[name] = lattices[name]
+
+    _sort_lattices(root if root is not None else latticejson["root"])
+    if keep_unused:
+        while len(lattices_set) > 0:
+            _sort_lattices(lattices_set.pop())
+    else:
+        for lattice in lattices_set:
+            warn(f"Discard unused lattice '{lattice}'.")
+    return lattices_sorted
+
+
+def remove_unused(latticejson, root=None, warn_unused=False):
+    """Remove unused objects starting from the `root` lattice. Also sorts lattices."""
+    if root is None:
+        root = latticejson["root"]
+    elements = latticejson["elements"]
+    lattices = latticejson["lattices"]
+    elements_set = set(elements)
+    lattices_set = set(lattices)
+    elements_new = {}
+    lattices_new = {}
+
+    def _remove_unused(name):
+        try:
+            elements_set.remove(name)
+        except KeyError:
+            pass
+        else:
+            elements_new[name] = elements[name]
+            return
+
+        try:
+            lattices_set.remove(name)
+        except KeyError:
+            pass
+        else:
+            lattice = lattices[name]
+            for child in lattice:
+                _remove_unused(child)
+            lattices_new[name] = lattices[name]
+
+    _remove_unused(root)
+    latticejson_new = latticejson.copy()
+    latticejson_new["root"] = root
+    latticejson_new["elements"] = elements_new
+    latticejson_new["lattices"] = lattices_new
+    if warn_unused:
+        for obj in chain(elements_set, lattices_set):
+            warn(f"Discard unused object '{obj}'.")
+    return latticejson_new