compatible to pyyaml>=6.0

yaml/pyyaml#576
Signed-off-by: quan <[email protected]>

aten/src/ATen/cwrap_parser.py

@@ -1,5 +1,10 @@
+# from importlib.abc import Loader
 import yaml
-
+try:
+    # use faster C loader if available
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
 # follows similar logic to cwrap, ignores !inc, and just looks for [[]]
 
 
@@ -15,7 +20,7 @@ def parse(filename):
                 in_declaration = True
             elif line == ']]':
                 in_declaration = False
-                declaration = yaml.load('\n'.join(declaration_lines))
+                declaration = yaml.load('\n'.join(declaration_lines),Loader=Loader)
                 declarations.append(declaration)
             elif in_declaration:
                 declaration_lines.append(line)

test/test_namedtuple_return_api.py

@@ -15,14 +15,20 @@
     'triangular_solve'
 }
 
+try:
+    # use faster C loader if available
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
+
 
 class TestNamedTupleAPI(unittest.TestCase):
 
     def test_native_functions_yaml(self):
         operators_found = set()
         regex = re.compile(r"^(\w*)\(")
         file = open(aten_native_yaml, 'r')
-        for f in yaml.load(file.read()):
+        for f in yaml.load(file.read(),Loader=Loader):
             f = f['func']
             ret = f.split('->')[1].strip()
             name = regex.findall(f)[0]

tools/clang_tidy.py

@@ -156,6 +156,12 @@ def run_shell_commands_in_parallel(commands):
     finally:
         os.unlink(f.name)
 
+try:
+    # use faster C loader if available
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
+
 
 def run_clang_tidy(options, line_filters, files):
     """Executes the actual clang-tidy command in the shell."""
@@ -167,7 +173,7 @@ def run_clang_tidy(options, line_filters, files):
 
         with open(options.config_file) as config:
             # Here we convert the YAML config file to a JSON blob.
-            command += ["-config", json.dumps(yaml.load(config))]
+            command += ["-config", json.dumps(yaml.load(config,Loader=Loader))]
     command += options.extra_args
 
     if line_filters:

tools/cwrap/cwrap.py

@@ -6,6 +6,12 @@
     BeforeAfterCall, ConstantArguments, ReturnArguments, GILRelease
 from ..shared import cwrap_common
 
+try:
+    # use faster C loader if available
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
+
 
 class cwrap(object):
     BASE_INDENT_SIZE = 6
@@ -88,7 +94,7 @@ def wrap_declarations(self, declarations):
                 in_declaration = True
             elif line == ']]':
                 in_declaration = False
-                declaration = yaml.load('\n'.join(declaration_lines))
+                declaration = yaml.load('\n'.join(declaration_lines),Loader=Loader)
                 cwrap_common.set_declaration_defaults(declaration)
 
                 # Pass declaration in a list - maybe some plugins want to add