diff options
Diffstat (limited to 'yarnhelper.py')
| -rw-r--r-- | yarnhelper.py | 20 |
1 files changed, 12 insertions, 8 deletions
diff --git a/yarnhelper.py b/yarnhelper.py index 8f9da19..4ed2c03 100644 --- a/yarnhelper.py +++ b/yarnhelper.py @@ -26,6 +26,7 @@ import requests import yaml +datadir = os.environ.get('DATADIR', '.') variables_filename = os.environ.get('VARIABLES', 'vars.yaml') @@ -34,6 +35,7 @@ class YarnHelper(object): def __init__(self): self._env = dict(os.environ) self._next_match = 1 + self._filename = os.path.join(datadir, variables_filename) self._variables = None # None means not loaded, otherwise dict def set_environment(self, env): @@ -46,17 +48,19 @@ class YarnHelper(object): self._next_match += 1 return self._env[name] - def get_variable(self, name): + def get_variable(self, name, default=None): if self._variables is None: self._variables = self._load_variables() - if name not in self._variables: - raise Error('no variable {}'.format(name)) - return self._variables[name] + assert self._variables is not None + return self._variables.get(name, default) def _load_variables(self): - if os.path.exists(variables_filename): - with open(variables_filename, 'r') as f: - return yaml.safe_load(f) + if os.path.exists(self._filename): + with open(self._filename, 'r') as f: + data = f.read() + if data: + f.seek(0) + return yaml.safe_load(f) return {} def set_variable(self, name, value): @@ -66,7 +70,7 @@ class YarnHelper(object): self._save_variables(self._variables) def _save_variables(self, variables): - with open(variables_filename, 'w') as f: + with open(self._filename, 'w') as f: yaml.safe_dump(variables, f) def construct_aliased_http_request( |