awx/awxkit/awxkit/yaml_file.py
Ryan Petrello 9616cc6f78
import awxkit
Co-authored-by: Christopher Wang <cwang@ansible.com>
Co-authored-by: Jake McDermott <jmcdermott@ansible.com>
Co-authored-by: Jim Ladd <jladd@redhat.com>
Co-authored-by: Elijah DeLee <kdelee@redhat.com>
Co-authored-by: Alan Rominger <arominge@redhat.com>
Co-authored-by: Yanis Guenane <yanis@guenane.org>
2019-08-08 22:12:31 -04:00

98 lines
3.0 KiB
Python

import os
import yaml
import glob
import logging
from py.path import local
log = logging.getLogger(__name__)
file_pattern_cache = {}
file_path_cache = {}
class Loader(yaml.FullLoader):
def __init__(self, stream):
self._root = os.path.split(stream.name)[0]
super(Loader, self).__init__(stream)
Loader.add_constructor('!include', Loader.include)
Loader.add_constructor('!import', Loader.include)
def include(self, node):
if isinstance(node, yaml.ScalarNode):
return self.extractFile(self.construct_scalar(node))
elif isinstance(node, yaml.SequenceNode):
result = []
for filename in self.construct_sequence(node):
result += self.extractFile(filename)
return result
elif isinstance(node, yaml.MappingNode):
result = {}
for k, v in self.construct_mapping(node).items():
result[k] = self.extractFile(v)[k]
return result
else:
log.error("unrecognised node type in !include statement")
raise yaml.constructor.ConstructorError
def extractFile(self, filename):
file_pattern = os.path.join(self._root, filename)
log.debug('Will attempt to extract schema from: {0}'.format(file_pattern))
if file_pattern in file_pattern_cache:
log.debug('File pattern cache hit: {0}'.format(file_pattern))
return file_pattern_cache[file_pattern]
data = dict()
for file_path in glob.glob(file_pattern):
file_path = os.path.abspath(file_path)
if file_path in file_path_cache:
log.debug('Schema cache hit: {0}'.format(file_path))
path_data = file_path_cache[file_path]
else:
log.debug('Loading schema from {0}'.format(file_path))
with open(file_path, 'r') as f:
path_data = yaml.load(f, Loader)
file_path_cache[file_path] = path_data
data.update(path_data)
file_pattern_cache[file_pattern] = data
return data
def load_file(filename):
"""Loads a YAML file from the given filename.
If the filename is omitted or None, attempts will be made to load it from
its normal location in the parent of the utils directory.
The awx_data dict loaded with this method supports value randomization,
thanks to the RandomizeValues class. See that class for possible options
Example usage in data.yaml (quotes are important!):
top_level:
list:
- "{random_str}"
- "{random_int}"
- "{random_uuid}"
random_thing: "{random_string:24}"
"""
if filename is None:
this_file = os.path.abspath(__file__)
path = local(this_file).new(basename='../data.yaml')
else:
path = local(filename)
if path.check():
fp = path.open()
# FIXME - support load_all()
return yaml.load(fp, Loader=Loader)
else:
msg = 'Unable to load data file at %s' % path
raise Exception(msg)