1 files changed, 155 insertions, 0 deletions
diff --git a/taskcluster/taskgraph/util/templates.py b/taskcluster/taskgraph/util/templates.py
new file mode 100644
index 000000000..97620fa75
--- /dev/null
+++ b/taskcluster/taskgraph/util/templates.py
@@ -0,0 +1,155 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+import os
+
+import pystache
+import yaml
+import copy
+
+# Key used in template inheritance...
+INHERITS_KEY = '$inherits'
+
+
+def merge_to(source, dest):
+    '''
+    Merge dict and arrays (override scalar values)
+
+    Keys from source override keys from dest, and elements from lists in source
+    are appended to lists in dest.
+
+    :param dict source: to copy from
+    :param dict dest: to copy to (modified in place)
+    '''
+
+    for key, value in source.items():
+        # Override mismatching or empty types
+        if type(value) != type(dest.get(key)):  # noqa
+            dest[key] = source[key]
+            continue
+
+        # Merge dict
+        if isinstance(value, dict):
+            merge_to(value, dest[key])
+            continue
+
+        if isinstance(value, list):
+            dest[key] = dest[key] + source[key]
+            continue
+
+        dest[key] = source[key]
+
+    return dest
+
+
+def merge(*objects):
+    '''
+    Merge the given objects, using the semantics described for merge_to, with
+    objects later in the list taking precedence.  From an inheritance
+    perspective, "parents" should be listed before "children".
+
+    Returns the result without modifying any arguments.
+    '''
+    if len(objects) == 1:
+        return copy.deepcopy(objects[0])
+    return merge_to(objects[-1], merge(*objects[:-1]))
+
+
+class TemplatesException(Exception):
+    pass
+
+
+class Templates():
+    '''
+    The taskcluster integration makes heavy use of yaml to describe tasks this
+    class handles the loading/rendering.
+    '''
+
+    def __init__(self, root):
+        '''
+        Initialize the template render.
+
+        :param str root: Root path where to load yaml files.
+        '''
+        if not root:
+            raise TemplatesException('Root is required')
+
+        if not os.path.isdir(root):
+            raise TemplatesException('Root must be a directory')
+
+        self.root = root
+
+    def _inherits(self, path, obj, properties, seen):
+        blueprint = obj.pop(INHERITS_KEY)
+        seen.add(path)
+
+        # Resolve the path here so we can detect circular references.
+        template = self.resolve_path(blueprint.get('from'))
+        variables = blueprint.get('variables', {})
+
+        # Passed parameters override anything in the task itself.
+        for key in properties:
+            variables[key] = properties[key]
+
+        if not template:
+            msg = '"{}" inheritance template missing'.format(path)
+            raise TemplatesException(msg)
+
+        if template in seen:
+            msg = 'Error while handling "{}" in "{}" circular template' + \
+                  'inheritance seen \n  {}'
+            raise TemplatesException(msg.format(path, template, seen))
+
+        try:
+            out = self.load(template, variables, seen)
+        except TemplatesException as e:
+            msg = 'Error expanding parent ("{}") of "{}" original error {}'
+            raise TemplatesException(msg.format(template, path, str(e)))
+
+        # Anything left in obj is merged into final results (and overrides)
+        return merge_to(obj, out)
+
+    def render(self, path, content, parameters, seen):
+        '''
+        Renders a given yaml string.
+
+        :param str path:  used to prevent infinite recursion in inheritance.
+        :param str content: Of yaml file.
+        :param dict parameters: For mustache templates.
+        :param set seen: Seen files (used for inheritance)
+        '''
+        content = pystache.render(content, parameters)
+        result = yaml.load(content)
+
+        # In addition to the usual template logic done by mustache we also
+        # handle special '$inherit' dict keys.
+        if isinstance(result, dict) and INHERITS_KEY in result:
+            return self._inherits(path, result, parameters, seen)
+
+        return result
+
+    def resolve_path(self, path):
+        return os.path.join(self.root, path)
+
+    def load(self, path, parameters=None, seen=None):
+        '''
+        Load an render the given yaml path.
+
+        :param str path: Location of yaml file to load (relative to root).
+        :param dict parameters: To template yaml file with.
+        '''
+        seen = seen or set()
+
+        if not path:
+            raise TemplatesException('path is required')
+
+        path = self.resolve_path(path)
+
+        if not os.path.isfile(path):
+            raise TemplatesException('"{}" is not a file'.format(path))
+
+        content = open(path).read()
+        return self.render(path, content, parameters, seen)