diff --git a/python/loadleveller/jobfile.py b/python/loadleveller/jobfile.py index b55112e683426ca466706c2b3c83e2c2509a3d88..eb37bb07d690c08625dab68002cc36883448b1d5 100644 --- a/python/loadleveller/jobfile.py +++ b/python/loadleveller/jobfile.py @@ -3,6 +3,11 @@ import os import subprocess import errno +try: + from yaml import CSafeLoader as SafeLoader +except ImportError: + from yaml import SafeLoader + '''Helpers for handling loadleveller jobfiles/scripts. For lack of a better idea, the job description files of loadleveller are actually executables that output a more verbose yaml parameter file to stdout. Use the taskmaker module to write the input scripts.''' class JobFileGenError(Exception): @@ -19,7 +24,7 @@ class JobFile: raise JobFileGenError('Generation script "{}" had a non-zero return code. Treating as error.'.format(filename)) try: - parsed_job = yaml.load(self.raw_jobfile, Loader=yaml.CSafeLoader) + parsed_job = yaml.load(self.raw_jobfile, Loader=SafeLoader) self.__dict__.update(parsed_job) except Exception as e: raise JobFileGenError('Could not parse job generation script output: {}'.format(e)) diff --git a/python/loadleveller/mcextract.py b/python/loadleveller/mcextract.py index 3927f23dbb64fc3e8527f6b76eb4a6f814e2049b..9c4c8fe5e9f8c9f8bdc1ed16adf0720f2334b832 100644 --- a/python/loadleveller/mcextract.py +++ b/python/loadleveller/mcextract.py @@ -2,6 +2,11 @@ import yaml import numpy as np import itertools +try: + from yaml import CSafeLoader as SafeLoader +except ImportError: + from yaml import SafeLoader + '''This module can be used to easily extract Monte Carlo results from the *.results.yml file produced by the loadleveller library.''' class Observable: @@ -16,7 +21,7 @@ class Observable: class MCArchive: def __init__(self, filename): with open(filename, 'r') as f: - doc = yaml.load(f, Loader=yaml.CSafeLoader) + doc = yaml.load(f, Loader=SafeLoader) param_names = set(sum([list(task['parameters'].keys()) for task in doc], [])) observable_names = set(sum([list(task['results'].keys()) for task in doc], [])) diff --git a/python/loadleveller/taskmaker.py b/python/loadleveller/taskmaker.py index f938ff3167cff270bbe5f5e2045a70b6f33f83c3..4e49fca50abfb67118649c1d772e1a32ad4683cb 100644 --- a/python/loadleveller/taskmaker.py +++ b/python/loadleveller/taskmaker.py @@ -3,12 +3,17 @@ import os import yaml import numpy +try: + from yaml import CSafeLoader as SafeLoader, CSafeDumper as SafeDumper +except ImportError: + from yaml import SafeLoader, SafeDumper + def _expand_path(path): return os.path.abspath(os.path.expandvars(os.path.expanduser(path))) def JobConfig(filename): with open(_expand_path(filename), 'r') as f: - return yaml.load(f, Loader=yaml.CSafeLoader) + return yaml.load(f, Loader=SafeLoader) class TaskMaker: def __init__(self, name, jobconfig): @@ -41,4 +46,4 @@ class TaskMaker: task_dict[k] = v jobfile_dict['tasks'][task_name] = task_dict - print(yaml.dump(jobfile_dict, Dumper=yaml.CSafeDumper)) + print(yaml.dump(jobfile_dict, Dumper=SafeDumper))