pytorch

1
#!/usr/bin/env python3
2
"""
3
Verify that it is possible to round-trip native_functions.yaml via ruamel under some
4
configuration.  Keeping native_functions.yaml consistent in this way allows us to
5
run codemods on the file using ruamel without introducing line noise.  Note that we don't
6
want to normalize the YAML file, as that would to lots of spurious lint failures.  Anything
7
that ruamel understands how to roundtrip, e.g., whitespace and comments, is OK!
8

9
ruamel is a bit picky about inconsistent indentation, so you will have to indent your
10
file properly.  Also, if you are working on changing the syntax of native_functions.yaml,
11
you may find that you want to use some format that is not what ruamel prefers.  If so,
12
it is OK to modify this script (instead of reformatting native_functions.yaml)--the point
13
is simply to make sure that there is *some* configuration of ruamel that can round trip
14
the YAML, not to be prescriptive about it.
15
"""
16

17
import difflib
18
import sys
19
from io import StringIO
20
from pathlib import Path
21

22
import ruamel.yaml  # type: ignore[import]
23

24

25
def fn(base: str) -> str:
26
    return str(base / Path("aten/src/ATen/native/native_functions.yaml"))
27

28

29
with open(Path(__file__).parent.parent.parent / fn(".")) as f:
30
    contents = f.read()
31

32
yaml = ruamel.yaml.YAML()  # type: ignore[attr-defined]
33
yaml.preserve_quotes = True  # type: ignore[assignment]
34
yaml.width = 1000  # type: ignore[assignment]
35
yaml.boolean_representation = ["False", "True"]  # type: ignore[attr-defined]
36
r = yaml.load(contents)
37

38
# Cuz ruamel's author intentionally didn't include conversion to string
39
# https://stackoverflow.com/questions/47614862/best-way-to-use-ruamel-yaml-to-dump-to-string-not-to-stream
40
string_stream = StringIO()
41
yaml.dump(r, string_stream)
42
new_contents = string_stream.getvalue()
43
string_stream.close()
44

45
if contents != new_contents:
46
    print(
47
        """\
48

49
## LINT FAILURE: native_functions.yaml ##
50

51
native_functions.yaml failed lint; please apply the diff below to fix lint.
52
If you think this is in error, please see .github/scripts/lint_native_functions.py
53
""",
54
        file=sys.stderr,
55
    )
56
    sys.stdout.writelines(
57
        difflib.unified_diff(
58
            contents.splitlines(True), new_contents.splitlines(True), fn("a"), fn("b")
59
        )
60
    )
61
    sys.exit(1)
62