Merge pull request github#11498 from github/redsun82/swift-codegen

Swift: enhance `codegen` UX

Author: redsun82

swift/codegen/codegen.py

@@ -42,6 +42,8 @@ def _parse_args() -> argparse.Namespace:
                    help="output directory for generated C++ files, required if trap or cpp is provided to --generate")
     p.add_argument("--generated-registry", type=_abspath, default=paths.swift_dir / "ql/.generated.list",
                    help="registry file containing information about checked-in generated code")
+    p.add_argument("--force", "-f", action="store_true",
+                   help="generate all files without skipping unchanged files and overwriting modified ones")
     return p.parse_args()
 
 

swift/codegen/generators/qlgen.py

@@ -302,7 +302,8 @@ def generate(opts, renderer):
 
     imports = {}
 
-    with renderer.manage(generated=generated, stubs=stubs, registry=opts.generated_registry) as renderer:
+    with renderer.manage(generated=generated, stubs=stubs, registry=opts.generated_registry,
+                         force=opts.force) as renderer:
 
         db_classes = [cls for cls in classes.values() if not cls.ipa]
         renderer.render(ql.DbClasses(db_classes), out / "Raw.qll")

swift/codegen/lib/render.py

@@ -59,8 +59,8 @@ def _do_write(self, mnemonic: str, contents: str, output: pathlib.Path):
         log.debug(f"{mnemonic}: generated {output.name}")
 
     def manage(self, generated: typing.Iterable[pathlib.Path], stubs: typing.Iterable[pathlib.Path],
-               registry: pathlib.Path) -> "RenderManager":
-        return RenderManager(self._swift_dir, generated, stubs, registry)
+               registry: pathlib.Path, force: bool = False) -> "RenderManager":
+        return RenderManager(self._swift_dir, generated, stubs, registry, force)
 
 
 class RenderManager(Renderer):
@@ -87,9 +87,10 @@ class Hashes:
 
     def __init__(self, swift_dir: pathlib.Path, generated: typing.Iterable[pathlib.Path],
                  stubs: typing.Iterable[pathlib.Path],
-                 registry: pathlib.Path):
+                 registry: pathlib.Path, force: bool = False):
         super().__init__(swift_dir)
         self._registry_path = registry
+        self._force = force
         self._hashes = {}
         self.written = set()
         self._existing = set()
@@ -103,12 +104,18 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        for f in self._existing - self._skipped - self.written:
-            self._hashes.pop(self._get_path(f), None)
-            f.unlink(missing_ok=True)
-            log.info(f"removed {f.name}")
-        for f in self.written:
-            self._hashes[self._get_path(f)].post = self._hash_file(f)
+        if exc_val is None:
+            for f in self._existing - self._skipped - self.written:
+                self._hashes.pop(self._get_path(f), None)
+                f.unlink(missing_ok=True)
+                log.info(f"removed {f.name}")
+            for f in self.written:
+                self._hashes[self._get_path(f)].post = self._hash_file(f)
+        else:
+            # if an error was encountered, drop already written files from the registry
+            # so that they get the chance to be regenerated again during the next run
+            for f in self.written:
+                self._hashes.pop(self._get_path(f), None)
         self._dump_registry()
 
     def _do_write(self, mnemonic: str, contents: str, output: pathlib.Path):
@@ -126,10 +133,13 @@ def _process_generated(self, generated: typing.Iterable[pathlib.Path]):
         for f in generated:
             self._existing.add(f)
             rel_path = self._get_path(f)
-            if rel_path not in self._hashes:
+            if self._force:
+                pass
+            elif rel_path not in self._hashes:
                 log.warning(f"{rel_path} marked as generated but absent from the registry")
             elif self._hashes[rel_path].post != self._hash_file(f):
-                raise Error(f"{rel_path} is generated but was modified, please revert the file")
+                raise Error(f"{rel_path} is generated but was modified, please revert the file "
+                            "or pass --force to overwrite")
 
     def _process_stubs(self, stubs: typing.Iterable[pathlib.Path]):
         for f in stubs:
@@ -138,10 +148,13 @@ def _process_stubs(self, stubs: typing.Iterable[pathlib.Path]):
                 self._hashes.pop(rel_path, None)
                 continue
             self._existing.add(f)
-            if rel_path not in self._hashes:
+            if self._force:
+                pass
+            elif rel_path not in self._hashes:
                 log.warning(f"{rel_path} marked as stub but absent from the registry")
             elif self._hashes[rel_path].post != self._hash_file(f):
-                raise Error(f"{rel_path} is a stub marked as generated, but it was modified")
+                raise Error(f"{rel_path} is a stub marked as generated, but it was modified, "
+                            "please remove the `// generated` header, revert the file or pass --force to overwrite it")
 
     @staticmethod
     def is_customized_stub(file: pathlib.Path) -> bool:
@@ -165,12 +178,18 @@ def _hash_string(data: str) -> str:
         return h.hexdigest()
 
     def _load_registry(self):
-        with open(self._registry_path) as reg:
-            for line in reg:
-                filename, prehash, posthash = line.split()
-                self._hashes[pathlib.Path(filename)] = self.Hashes(prehash, posthash)
+        if self._force:
+            return
+        try:
+            with open(self._registry_path) as reg:
+                for line in reg:
+                    filename, prehash, posthash = line.split()
+                    self._hashes[pathlib.Path(filename)] = self.Hashes(prehash, posthash)
+        except FileNotFoundError:
+            pass
 
     def _dump_registry(self):
+        self._registry_path.parent.mkdir(parents=True, exist_ok=True)
         with open(self._registry_path, 'w') as out:
             for f, hashes in sorted(self._hashes.items()):
                 print(f, hashes.pre, hashes.post, file=out)

swift/codegen/test/test_qlgen.py

@@ -48,6 +48,7 @@ def qlgen_opts(opts):
     opts.generated_registry = generated_registry_path()
     opts.ql_format = True
     opts.swift_dir = paths.swift_dir
+    opts.force = False
     return opts
 
 
@@ -430,7 +431,9 @@ def test_format_error(opts, generate, render_manager, run_mock):
         generate([schema.Class('A')])
 
 
-def test_manage_parameters(opts, generate, renderer):
+@pytest.mark.parametrize("force", [False, True])
+def test_manage_parameters(opts, generate, renderer, force):
+    opts.force = force
     ql_a = opts.ql_output / "A.qll"
     ql_b = opts.ql_output / "B.qll"
     stub_a = opts.ql_stub_output / "A.qll"
@@ -448,7 +451,7 @@ def test_manage_parameters(opts, generate, renderer):
     generate([schema.Class('A')])
     assert renderer.mock_calls == [
         mock.call.manage(generated={ql_a, ql_b, test_a, test_b, import_file()}, stubs={stub_a, stub_b},
-                         registry=opts.generated_registry)
+                         registry=opts.generated_registry, force=force)
     ]
 
 

swift/codegen/test/test_render.py

@@ -76,6 +76,24 @@ def test_managed_render(pystache_renderer, sut):
     ]
 
 
+def test_managed_render_with_no_registry(pystache_renderer, sut):
+    data = mock.Mock(spec=("template",))
+    text = "some text"
+    pystache_renderer.render_name.side_effect = (text,)
+    output = paths.swift_dir / "some/output.txt"
+    registry = paths.swift_dir / "a/registry.list"
+
+    with sut.manage(generated=(), stubs=(), registry=registry) as renderer:
+        renderer.render(data, output)
+        assert renderer.written == {output}
+        assert_file(output, text)
+
+    assert_file(registry, f"some/output.txt {hash(text)} {hash(text)}\n")
+    assert pystache_renderer.mock_calls == [
+        mock.call.render_name(data.template, data, generator=paths.exe_file.relative_to(paths.swift_dir)),
+    ]
+
+
 def test_managed_render_with_post_processing(pystache_renderer, sut):
     data = mock.Mock(spec=("template",))
     text = "some text"
@@ -192,6 +210,45 @@ def test_managed_render_with_modified_stub_file_not_marked_as_generated(pystache
     assert_file(registry, "")
 
 
+class MyError(Exception):
+    pass
+
+
+def test_managed_render_exception_drops_written_from_registry(pystache_renderer, sut):
+    data = mock.Mock(spec=("template",))
+    text = "some text"
+    pystache_renderer.render_name.side_effect = (text,)
+    output = paths.swift_dir / "some/output.txt"
+    registry = paths.swift_dir / "a/registry.list"
+    write(output, text)
+    write(registry, "a a a\n"
+                    f"some/output.txt whatever {hash(text)}\n"
+                    "b b b")
+
+    with pytest.raises(MyError):
+        with sut.manage(generated=(), stubs=(), registry=registry) as renderer:
+            renderer.render(data, output)
+            raise MyError
+
+    assert_file(registry, "a a a\nb b b\n")
+
+
+def test_managed_render_exception_does_not_erase(pystache_renderer, sut):
+    output = paths.swift_dir / "some/output.txt"
+    stub = paths.swift_dir / "some/stub.txt"
+    registry = paths.swift_dir / "a/registry.list"
+    write(output)
+    write(stub, "// generated bla bla")
+    write(registry)
+
+    with pytest.raises(MyError):
+        with sut.manage(generated=(output,), stubs=(stub,), registry=registry) as renderer:
+            raise MyError
+
+    assert output.is_file()
+    assert stub.is_file()
+
+
 def test_render_with_extensions(pystache_renderer, sut):
     data = mock.Mock(spec=("template", "extensions"))
     data.template = "test_template"
@@ -210,5 +267,70 @@ def test_render_with_extensions(pystache_renderer, sut):
         assert_file(expected_output, expected_contents)
 
 
+def test_managed_render_with_force_not_skipping_generated_file(pystache_renderer, sut):
+    data = mock.Mock(spec=("template",))
+    output = paths.swift_dir / "some/output.txt"
+    some_output = "some output"
+    registry = paths.swift_dir / "a/registry.list"
+    write(output, some_output)
+    write(registry, f"some/output.txt {hash(some_output)} {hash(some_output)}\n")
+
+    pystache_renderer.render_name.side_effect = (some_output,)
+
+    with sut.manage(generated=(output,), stubs=(), registry=registry, force=True) as renderer:
+        renderer.render(data, output)
+        assert renderer.written == {output}
+        assert_file(output, some_output)
+
+    assert_file(registry, f"some/output.txt {hash(some_output)} {hash(some_output)}\n")
+    assert pystache_renderer.mock_calls == [
+        mock.call.render_name(data.template, data, generator=paths.exe_file.relative_to(paths.swift_dir)),
+    ]
+
+
+def test_managed_render_with_force_not_skipping_stub_file(pystache_renderer, sut):
+    data = mock.Mock(spec=("template",))
+    stub = paths.swift_dir / "some/stub.txt"
+    some_output = "// generated some output"
+    some_processed_output = "// generated some processed output"
+    registry = paths.swift_dir / "a/registry.list"
+    write(stub, some_processed_output)
+    write(registry, f"some/stub.txt {hash(some_output)} {hash(some_processed_output)}\n")
+
+    pystache_renderer.render_name.side_effect = (some_output,)
+
+    with sut.manage(generated=(), stubs=(stub,), registry=registry, force=True) as renderer:
+        renderer.render(data, stub)
+        assert renderer.written == {stub}
+        assert_file(stub, some_output)
+
+    assert_file(registry, f"some/stub.txt {hash(some_output)} {hash(some_output)}\n")
+    assert pystache_renderer.mock_calls == [
+        mock.call.render_name(data.template, data, generator=paths.exe_file.relative_to(paths.swift_dir)),
+    ]
+
+
+def test_managed_render_with_force_ignores_modified_generated_file(sut):
+    output = paths.swift_dir / "some/output.txt"
+    some_processed_output = "// some processed output"
+    registry = paths.swift_dir / "a/registry.list"
+    write(output, "// something else")
+    write(registry, f"some/output.txt whatever {hash(some_processed_output)}\n")
+
+    with sut.manage(generated=(output,), stubs=(), registry=registry, force=True):
+        pass
+
+
+def test_managed_render_with_force_ignores_modified_stub_file_still_marked_as_generated(sut):
+    stub = paths.swift_dir / "some/stub.txt"
+    some_processed_output = "// generated some processed output"
+    registry = paths.swift_dir / "a/registry.list"
+    write(stub, "// generated something else")
+    write(registry, f"some/stub.txt whatever {hash(some_processed_output)}\n")
+
+    with sut.manage(generated=(), stubs=(stub,), registry=registry, force=True):
+        pass
+
+
 if __name__ == '__main__':
     sys.exit(pytest.main([__file__] + sys.argv[1:]))