Implement `@snapshot` check for htmldocck

This form of check allows performing snapshot tests (à la `src/test/ui`)
on rustdoc HTML output, making it easier to create and update tests.

See this Zulip thread [1] for more information about the motivation for
this change.

[1]: https://zulip-archive.rust-lang.org/stream/266220-rustdoc/topic/HTML.20snapshot.20tests.html#262651142

2 files changed, 82 insertions, 10 deletions

diff --git a/src/etc/htmldocck.py b/src/etc/htmldocck.py
index 8647db5a45d..48a341ffe08 100644
--- a/src/etc/htmldocck.py
+++ b/src/etc/htmldocck.py
@@ -90,10 +90,20 @@ There are a number of supported commands:
   highlights for example. If you want to simply check for the presence of
   a given node or attribute, use an empty string (`""`) as a `PATTERN`.
 
-* `@count PATH XPATH COUNT' checks for the occurrence of the given XPath
+* `@count PATH XPATH COUNT` checks for the occurrence of the given XPath
   in the specified file. The number of occurrences must match the given
   count.
 
+* `@snapshot NAME PATH XPATH` creates a snapshot test named NAME.
+  A snapshot test captures a subtree of the DOM, at the location
+  determined by the XPath, and compares it to a pre-recorded value
+  in a file. The file's name is the test's name with the `.rs` extension
+  replaced with `.NAME.html`, where NAME is the snapshot's name.
+
+  htmldocck supports the `--bless` option to accept the current subtree
+  as expected, saving it to the file determined by the snapshot's name.
+  compiletest's `--bless` flag is forwarded to htmldocck.
+
 * `@has-dir PATH` checks for the existence of the given directory.
 
 All conditions can be negated with `!`. `@!has foo/type.NoSuch.html`
@@ -137,6 +147,10 @@ except NameError:
 
 channel = os.environ["DOC_RUST_LANG_ORG_CHANNEL"]
 
+# Initialized in main
+rust_test_path = None
+bless = None
+
 class CustomHTMLParser(HTMLParser):
     """simplified HTML parser.
 
@@ -387,6 +401,32 @@ def get_tree_count(tree, path):
     return len(tree.findall(path))
 
 
+def check_snapshot(snapshot_name, tree):
+    assert rust_test_path.endswith('.rs')
+    snapshot_path = '{}.{}.{}'.format(rust_test_path[:-3], snapshot_name, 'html')
+    try:
+        with open(snapshot_path, 'r') as snapshot_file:
+            expected_str = snapshot_file.read()
+    except FileNotFoundError:
+        if bless:
+            expected_str = None
+        else:
+            raise FailedCheck('No saved snapshot value')
+
+    actual_str = ET.tostring(tree).decode('utf-8')
+
+    if expected_str != actual_str:
+        if bless:
+            with open(snapshot_path, 'w') as snapshot_file:
+                snapshot_file.write(actual_str)
+        else:
+            print('--- expected ---\n')
+            print(expected_str)
+            print('\n\n--- actual ---\n')
+            print(actual_str)
+            print()
+            raise FailedCheck('Actual snapshot value is different than expected')
+
 def stderr(*args):
     if sys.version_info.major < 3:
         file = codecs.getwriter('utf-8')(sys.stderr)
@@ -448,6 +488,28 @@ def check_command(c, cache):
                 ret = expected == found
             else:
                 raise InvalidCheck('Invalid number of @{} arguments'.format(c.cmd))
+
+        elif c.cmd == 'snapshot':  # snapshot test
+            if len(c.args) == 3:  # @snapshot <snapshot-name> <html-path> <xpath>
+                [snapshot_name, html_path, pattern] = c.args
+                tree = cache.get_tree(html_path)
+                xpath = normalize_xpath(pattern)
+                subtrees = tree.findall(xpath)
+                if len(subtrees) == 1:
+                    [subtree] = subtrees
+                    try:
+                        check_snapshot(snapshot_name, subtree)
+                        ret = True
+                    except FailedCheck as err:
+                        cerr = str(err)
+                        ret = False
+                elif len(subtrees) == 0:
+                    raise FailedCheck('XPATH did not match')
+                else:
+                    raise FailedCheck('Expected 1 match, but found {}'.format(len(subtrees)))
+            else:
+                raise InvalidCheck('Invalid number of @{} arguments'.format(c.cmd))
+
         elif c.cmd == 'has-dir':  # has-dir test
             if len(c.args) == 1:  # @has-dir <path> = has-dir test
                 try:
@@ -458,11 +520,13 @@ def check_command(c, cache):
                     ret = False
             else:
                 raise InvalidCheck('Invalid number of @{} arguments'.format(c.cmd))
+
         elif c.cmd == 'valid-html':
             raise InvalidCheck('Unimplemented @valid-html')
 
         elif c.cmd == 'valid-links':
             raise InvalidCheck('Unimplemented @valid-links')
+
         else:
             raise InvalidCheck('Unrecognized @{}'.format(c.cmd))
 
@@ -483,11 +547,19 @@ def check(target, commands):
 
 
 if __name__ == '__main__':
-    if len(sys.argv) != 3:
-        stderr('Usage: {} <doc dir> <template>'.format(sys.argv[0]))
+    if len(sys.argv) not in [3, 4]:
+        stderr('Usage: {} <doc dir> <template> [--bless]'.format(sys.argv[0]))
         raise SystemExit(1)
 
-    check(sys.argv[1], get_commands(sys.argv[2]))
+    rust_test_path = sys.argv[2]
+    if len(sys.argv) > 3 and sys.argv[3] == '--bless':
+        bless = True
+    else:
+        # We only support `--bless` at the end of the arguments.
+        # This assert is to prevent silent failures.
+        assert '--bless' not in sys.argv
+        bless = False
+    check(sys.argv[1], get_commands(rust_test_path))
     if ERR_COUNT:
         stderr("\nEncountered {} errors".format(ERR_COUNT))
         raise SystemExit(1)
diff --git a/src/tools/compiletest/src/runtest.rs b/src/tools/compiletest/src/runtest.rs
index 3c85b9076dd..d0ad4b80cde 100644
--- a/src/tools/compiletest/src/runtest.rs
+++ b/src/tools/compiletest/src/runtest.rs
@@ -2217,12 +2217,12 @@ impl<'test> TestCx<'test> {
             self.check_rustdoc_test_option(proc_res);
         } else {
             let root = self.config.find_rust_src_root().unwrap();
-            let res = self.cmd2procres(
-                Command::new(&self.config.docck_python)
-                    .arg(root.join("src/etc/htmldocck.py"))
-                    .arg(&out_dir)
-                    .arg(&self.testpaths.file),
-            );
+            let mut cmd = Command::new(&self.config.docck_python);
+            cmd.arg(root.join("src/etc/htmldocck.py")).arg(&out_dir).arg(&self.testpaths.file);
+            if self.config.bless {
+                cmd.arg("--bless");
+            }
+            let res = self.cmd2procres(&mut cmd);
             if !res.status.success() {
                 self.fatal_proc_rec_with_ctx("htmldocck failed!", &res, |mut this| {
                     this.compare_to_default_rustdoc(&out_dir)