Better autoformat (#26)

* autoformat multiple files at once
* add newline at end of file

Author: felix-andreas

latticejson/cli.py

@@ -1,6 +1,7 @@
 import click
 import json
 from pathlib import Path
+import itertools
 
 from .validate import validate_file
 from .io import convert as _convert
@@ -10,7 +11,7 @@
 
 
 FORMATS = "json", "lte"
-print_latticejson = lambda obj: print(json.dumps(obj, cls=CompactJSONEncoder, indent=4))
+dump_latticejson = lambda obj: json.dumps(obj, cls=CompactJSONEncoder, indent=4)
 
 
 @click.group()
@@ -39,7 +40,7 @@ def convert(file, from_, to):
     if from_ is None:
         from_ = path.suffix[1:]
 
-    print(_convert(path.read_text(), from_, to))
+    click.echo(_convert(path.read_text(), from_, to))
 
 
 @main.command()
@@ -54,15 +55,29 @@ def validate(file):
 def parse_elegant(file):
     """Parse elegant file but do not convert to LatticeJSON."""
     text = Path(file).read_text()
-    print(json.dumps(_parse_elegant(text), cls=CompactJSONEncoder, indent=4))
+    click.echo(dump_latticejson(_parse_elegant(text)))
 
 
 @main.command()
-@click.argument("file", type=click.Path(exists=True))
-def autoformat(file):
+@click.argument("files", nargs=-1, type=click.Path(exists=True))
+@click.option(
+    "--dry-run",
+    "-d",
+    is_flag=True,
+    help="Don't write the files back, just output the formatted files.",
+)
+def autoformat(files, dry_run):
     """Format a LatticeJSON file."""
-    latticejson = json.loads(Path(file).read_text())
-    print(json.dumps(latticejson, cls=CompactJSONEncoder, indent=4))
+    for path in itertools.chain.from_iterable(
+        path.rglob("*.json") if path.is_dir() else (path,) for path in map(Path, files)
+    ):
+        latticejson = json.loads(path.read_text())
+        formatted = dump_latticejson(latticejson)
+        click.secho(f"reformatted {path}", bold=True)
+        if dry_run:
+            click.echo(formatted)
+        else:
+            path.write_text(formatted)
 
 
 @main.command()
@@ -74,5 +89,5 @@ def migrate(file, from_, to):
     text = Path(file).read_text()
     initial_version = from_.split(".")
     final_version = to.split(".")
-    res = _migrate(json.loads(text), initial_version, final_version)
-    print(json.dumps(res, cls=CompactJSONEncoder, indent=4))
+    latticejson = _migrate(json.loads(text), initial_version, final_version)
+    click.echo(dump_latticejson(latticejson))

latticejson/format.py

@@ -2,17 +2,17 @@
 
 
 class CompactJSONEncoder(json.JSONEncoder):
-    """A JSON Encoder which only indents the first two level."""
+    """A JSON Encoder which only indents the first two levels."""
 
     def encode(self, obj, level=0):
-        if isinstance(obj, (list, tuple)):
-            return f"[{', '.join(json.dumps(el) for el in obj)}]"
-        elif isinstance(obj, dict) and level < 2:
-            indent = (level + 1) * self.indent * " "
-            items = ",\n".join(
-                f"{indent}{json.dumps(key)}: {self.encode(value, level=level+1)}"
+        if isinstance(obj, dict) and level < 2:
+            items_indent = (level + 1) * self.indent * " "
+            items_string = ",\n".join(
+                f"{items_indent}{json.dumps(key)}: {self.encode(value, level=level+1)}"
                 for key, value in obj.items()
             )
-            return f"{{\n{items}\n{level * self.indent * ' '}}}"
+            dict_indent = level * self.indent * " "
+            newline = "\n" if level == 0 else ""
+            return f"{{\n{items_string}\n{dict_indent}}}{newline}"
         else:
             return json.dumps(obj)