diff --git a/aten/src/ATen/cwrap_parser.py b/aten/src/ATen/cwrap_parser.py index f020dd030..880a69bd3 100644 --- a/aten/src/ATen/cwrap_parser.py +++ b/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) diff --git a/test/test_namedtuple_return_api.py b/test/test_namedtuple_return_api.py index b5471764a..0b89b7b8a 100644 --- a/test/test_namedtuple_return_api.py +++ b/test/test_namedtuple_return_api.py @@ -15,6 +15,12 @@ '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): @@ -22,7 +28,7 @@ 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] diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index edbd3d796..fba8bcbaa 100644 --- a/tools/clang_tidy.py +++ b/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: diff --git a/tools/cwrap/cwrap.py b/tools/cwrap/cwrap.py index 1444c0864..7316034d0 100644 --- a/tools/cwrap/cwrap.py +++ b/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