Update scripts/compile_tests/update_failures.py (pytorch#120529)

In order to unbreak this script, I have only tested with
```
./scripts/compile_tests/update_failures.py 97918e8
```

Pull Request resolved: pytorch#120529
Approved by: https://github.com/zou3519

scripts/compile_tests/update_failures.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 import argparse
+import os
 import pathlib
+import subprocess
 
 from common import (
     download_reports,
@@ -13,9 +15,9 @@
 )
 
 """
-Usage: update_failures.py /path/to/dynamo_test_failures.py commit_sha
+Usage: update_failures.py /path/to/dynamo_test_failures.py /path/to/test commit_sha
 
-Best-effort updates the xfail and skip lists in dynamo_test_failures.py
+Best-effort updates the xfail and skip files under test directory
 by parsing test reports.
 
 You'll need to provide the commit_sha for the latest commit on a PR
@@ -35,19 +37,28 @@
 """
 
 
-def patch_file(filename, unexpected_successes, new_xfails, new_skips, unexpected_skips):
-    with open(filename, "r") as f:
-        text = f.readlines()
+def patch_file(
+    filename, test_dir, unexpected_successes, new_xfails, new_skips, unexpected_skips
+):
+    failures_directory = os.path.join(test_dir, "dynamo_expected_failures")
+    skips_directory = os.path.join(test_dir, "dynamo_skips")
 
-    new_text = []
+    dynamo_expected_failures = set(os.listdir(failures_directory))
+    dynamo_skips = set(os.listdir(skips_directory))
 
-    i = 0
-    while True:
-        line = text[i]
-        if line.startswith("dynamo_expected_failures"):
-            break
-        new_text.append(line)
-        i += 1
+    # These are hand written skips
+    extra_dynamo_skips = set()
+    with open(filename, "r") as f:
+        start = False
+        for text in f.readlines():
+            text = text.strip()
+            if start:
+                if text == "}":
+                    break
+                extra_dynamo_skips.add(text.strip(',"'))
+            else:
+                if text == "extra_dynamo_skips = {":
+                    start = True
 
     def format(testcase):
         classname = testcase.attrib["classname"]
@@ -60,41 +71,30 @@ def format(testcase):
     formatted_unexpected_skips = {
         f"{format(test)}" for test in unexpected_skips.values()
     }
-    formatted_new_xfails = [
-        f'    "{format(test)}",  # {test.attrib["file"]}\n'
-        for test in new_xfails.values()
-    ]
-    formatted_new_skips = [
-        f'    "{format(test)}",  # {test.attrib["file"]}\n'
-        for test in new_skips.values()
-    ]
-
-    def is_in(lst, line):
-        splits = line.split('"')
-        if len(splits) < 3:
-            return None
-        test_name = splits[1]
-        if test_name in lst:
-            return test_name
-        return None
-
-    covered_unexpected_successes = set({})
+    formatted_new_xfails = [f"{format(test)}" for test in new_xfails.values()]
+    formatted_new_skips = [f"{format(test)}" for test in new_skips.values()]
+
+    def remove_file(path, name):
+        file = os.path.join(path, name)
+        cmd = ["git", "rm", file]
+        subprocess.run(cmd)
+
+    def add_file(path, name):
+        file = os.path.join(path, name)
+        with open(file, "w") as fp:
+            pass
+        cmd = ["git", "add", file]
+        subprocess.run(cmd)
+
+    covered_unexpected_successes = set()
 
     # dynamo_expected_failures
-    while True:
-        line = text[i]
-        match = is_in(formatted_unexpected_successes, line)
-        if match is not None:
-            covered_unexpected_successes.add(match)
-            i += 1
-            continue
-        if line == "}\n":
-            new_text.extend(formatted_new_xfails)
-            new_text.append(line)
-            i += 1
-            break
-        new_text.append(line)
-        i += 1
+    for test in dynamo_expected_failures:
+        if test in formatted_unexpected_successes:
+            covered_unexpected_successes.add(test)
+            remove_file(failures_directory, test)
+    for test in formatted_new_xfails:
+        add_file(failures_directory, test)
 
     leftover_unexpected_successes = (
         formatted_unexpected_successes - covered_unexpected_successes
@@ -108,29 +108,16 @@ def is_in(lst, line):
             print(stuff)
 
     # dynamo_skips
-    while True:
-        line = text[i]
-        match = is_in(formatted_unexpected_skips, line)
-        if match is not None:
-            i += 1
-            continue
-        if line == "}\n":
-            new_text.extend(formatted_new_skips)
-            break
-        if line == "dynamo_skips = {}\n":
-            new_text.extend("dynamo_skips = {\n")
-            new_text.extend(new_skips)
-            new_text.extend("}\n")
-            i += 1
-            break
-        new_text.append(line)
-        i += 1
-
-    for j in range(i, len(text)):
-        new_text.append(text[j])
-
-    with open(filename, "w") as f:
-        f.writelines(new_text)
+    for test in dynamo_skips:
+        if test in formatted_unexpected_skips:
+            remove_file(skips_directory, test)
+    for test in extra_dynamo_skips:
+        if test in formatted_unexpected_skips:
+            print(
+                f"WARNING: {test} in dynamo_test_failures.py needs to be removed manually"
+            )
+    for test in formatted_new_skips:
+        add_file(skips_directory, test)
 
 
 def get_intersection_and_outside(a_dict, b_dict):
@@ -151,7 +138,7 @@ def build_dict(keys):
     return build_dict(intersection), build_dict(outside)
 
 
-def update(filename, py38_dir, py311_dir, also_remove_skips):
+def update(filename, test_dir, py38_dir, py311_dir, also_remove_skips):
     def read_test_results(directory):
         xmls = open_test_results(directory)
         testcases = get_testcases(xmls)
@@ -192,7 +179,7 @@ def read_test_results(directory):
         f"{len(xfails)} new xfails, {len(all_skips)} new skips, {len(unexpected_skips)} new unexpected skips"
     )
     return patch_file(
-        filename, unexpected_successes, xfails, all_skips, unexpected_skips
+        filename, test_dir, unexpected_successes, xfails, all_skips, unexpected_skips
     )
 
 
@@ -211,6 +198,13 @@ def read_test_results(directory):
         ),
         help="Optional path to dynamo_test_failures.py",
     )
+    # test path
+    parser.add_argument(
+        "test_dir",
+        nargs="?",
+        default=str(pathlib.Path(__file__).absolute().parent.parent.parent / "test"),
+        help="Optional path to test folder",
+    )
     parser.add_argument(
         "commit",
         help=(
@@ -225,5 +219,6 @@ def read_test_results(directory):
     )
     args = parser.parse_args()
     assert pathlib.Path(args.filename).exists(), args.filename
+    assert pathlib.Path(args.test_dir).exists(), args.test_dir
     dynamo38, dynamo311 = download_reports(args.commit, ("dynamo38", "dynamo311"))
-    update(args.filename, dynamo38, dynamo311, args.also_remove_skips)
+    update(args.filename, args.test_dir, dynamo38, dynamo311, args.also_remove_skips)