From 90e39ca4fac71fbf6484aa7b96bce26e616c64ae Mon Sep 17 00:00:00 2001
From: Jennifer Taylor <jetaylor@fb.com>
Date: Tue, 27 Aug 2019 14:08:13 -0700
Subject: [PATCH] Add a tox environment for running codegen.

---
 libcst/codegen/generate.py                 | 102 +++++++++++++++++++++
 libcst/codegen/tests/test_codegen_clean.py |  15 +--
 tox.ini                                    |   7 ++
 3 files changed, 117 insertions(+), 7 deletions(-)
 create mode 100644 libcst/codegen/generate.py

diff --git a/libcst/codegen/generate.py b/libcst/codegen/generate.py
new file mode 100644
index 00000000..f452d83a
--- /dev/null
+++ b/libcst/codegen/generate.py
@@ -0,0 +1,102 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Usage:
+#
+# python -m libcst.codegen.generate --help
+# python -m libcst.codegen.generate visitors
+
+# pyre-strict
+import argparse
+import os
+import os.path
+import shutil
+import subprocess
+import sys
+from typing import List
+
+
+def format_file(fname: str) -> None:
+    with open(os.devnull, "w") as devnull:
+        subprocess.check_call(
+            ["isort", "-y", "-q", fname], stdout=devnull, stderr=devnull
+        )
+        subprocess.check_call(["black", fname], stdout=devnull, stderr=devnull)
+
+
+def codegen_visitors() -> None:
+    # First, back up the original file, since we have a nasty bootstrap problem.
+    # We're in a situation where we want to import libcst in order to get the
+    # valid nodes for visitors, but doing so means that we depend on ourselves.
+    # So, this attempts to keep the repo in a working state for as many operations
+    # as possible.
+    base = os.path.abspath(
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "../")
+    )
+    visitors_file = os.path.join(base, "_typed_visitor.py")
+    shutil.copyfile(visitors_file, f"{visitors_file}.bak")
+
+    try:
+        # Now that we backed up the file, lets codegen a new version.
+        # We import now, because this script does work on import.
+        import libcst.codegen.gen_visitor_functions as visitor_codegen
+
+        new_code = "\n".join(visitor_codegen.generated_code)
+        with open(visitors_file, "w") as fp:
+            fp.write(new_code)
+            fp.close()
+
+        # Now, see if the file we generated causes any import errors
+        # by attempting to run codegen again in a new process.
+        with open(os.devnull, "w") as devnull:
+            subprocess.check_call(
+                ["python3", "-m", "libcst.codegen.gen_visitor_functions"],
+                cwd=base,
+                stdout=devnull,
+            )
+
+        # If it worked, lets format the file
+        format_file(visitors_file)
+
+        # Since we were successful with importing, we can remove the backup.
+        os.remove(f"{visitors_file}.bak")
+
+        # Inform the user
+        print(f"Successfully generated a new {visitors_file} file.")
+    except Exception:
+        # On failure, we put the original file back, and keep the failed version
+        # for developers to look at.
+        print(
+            f"Failed to generated a new {visitors_file} file, failure "
+            + f"is saved in {visitors_file}.failed_generate.",
+            file=sys.stderr,
+        )
+        os.rename(visitors_file, f"{visitors_file}.failed_generate")
+        os.rename(f"{visitors_file}.bak", visitors_file)
+
+        # Reraise so we can debug
+        raise
+
+
+def main(cli_args: List[str]) -> int:
+    # Parse out arguments, run codegen
+    parser = argparse.ArgumentParser(description="Generate code for libcst.")
+    parser.add_argument(
+        "system",
+        metavar="SYSTEM",
+        help='System to generate code for. Valid values include: "visitors"',
+        type=str,
+    )
+    args = parser.parse_args(cli_args)
+    if args.system == "visitors":
+        codegen_visitors()
+        return 0
+    else:
+        print(f'Invalid system "{args.system}".')
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
diff --git a/libcst/codegen/tests/test_codegen_clean.py b/libcst/codegen/tests/test_codegen_clean.py
index c1eabbb2..d64951a4 100644
--- a/libcst/codegen/tests/test_codegen_clean.py
+++ b/libcst/codegen/tests/test_codegen_clean.py
@@ -6,9 +6,9 @@
 # pyre-strict
 import os
 import os.path
-import subprocess
 
 import libcst.codegen.gen_visitor_functions as visitor_codegen
+from libcst.codegen.generate import format_file
 from libcst.testing.utils import UnitTest
 
 
@@ -25,11 +25,12 @@ class TestCodegenClean(UnitTest):
         )
         with open(new_file, "w") as fp:
             fp.write(new_code)
-        with open(os.devnull, "w") as devnull:
-            subprocess.check_call(
-                ["isort", "-y", "-q", new_file], stdout=devnull, stderr=devnull
-            )
-            subprocess.check_call(["black", new_file], stdout=devnull, stderr=devnull)
+        try:
+            format_file(new_file)
+        except Exception:
+            # We failed to format, but this is probably due to invalid code that
+            # black doesn't like. This test will still fail and report to run codegen.
+            pass
         with open(new_file, "r") as fp:
             new_code = fp.read()
         os.remove(new_file)
@@ -43,5 +44,5 @@ class TestCodegenClean(UnitTest):
 
         # Now that we've done simple codegen, verify that it matches.
         self.assertTrue(
-            old_code == new_code, "libcst._typed_visitor needs to new codegen!"
+            old_code == new_code, "libcst._typed_visitor needs new codegen!"
         )
diff --git a/tox.ini b/tox.ini
index 0747611c..e3b74cf4 100644
--- a/tox.ini
+++ b/tox.ini
@@ -53,3 +53,10 @@ setenv =
     HYPOTHESIS = 1
 commands =
     python -m unittest libcst/tests/test_fuzz.py
+
+[testenv:codegen]
+deps =
+    -rrequirements.txt
+    -rrequirements-dev.txt
+commands =
+    python3 -m libcst.codegen.generate visitors