--- a/taskcluster/docs/try.rst
+++ b/taskcluster/docs/try.rst
@@ -101,42 +101,33 @@ from the ``try_task_config.json`` like t
     }
 
 Each key in the templates object denotes a new template to apply, and the value
 denotes extra context to use while rendering. When specified, a template will
 be applied to every task no matter what. If the template should only be applied
 to certain kinds of tasks, this needs to be specified in the template itself
 using JSON-e `condition statements`_.
 
-The context available to the JSON-e render aims to match that of ``actions``.
-It looks like this:
+The context available to the JSON-e render contains attributes from the
+:py:class:`taskgraph.task.Task` class. It looks like this:
 
 .. parsed-literal::
 
     {
-      "task": {
-        "payload": {
-          "env": { ... },
-          ...
-        }
-        "extra": {
-          "treeherder": { ... },
-          ...
-        },
-        "tags": { "kind": "<kind>", ... },
-        ...
-      },
-      "input": {
-        "enabled": 1,
-        ...
-      },
-      "taskId": "<task id>"
+      "attributes": task.attributes,
+      "kind": task.kind,
+      "task": task.task,
+      "taskId": task.task_id,
+      "input": ...
     }
 
-See the `existing templates`_ for examples.
+The ``input`` context can be any arbitrary value or object. What it contains
+depends on each specific template. Templates must return objects that have have
+either ``attributes`` or ``task`` as a top level key. All other top level keys
+will be ignored. See the `existing templates`_ for examples.
 
 Empty Try
 :::::::::
 
 If there is no try syntax or ``try_task_config.json``, the ``try_mode``
 parameter is None and no tasks are selected to run.  The resulting push will
 only have a decision task, but one with an "add jobs" action that can be used
 to add the desired jobs to the try push.

--- a/taskcluster/moz.build
+++ b/taskcluster/moz.build
@@ -22,8 +22,9 @@ with Files('scripts/**'):
 
 with Files('taskgraph/**'):
     BUG_COMPONENT = ('Taskcluster', 'Task Configuration')
 
 PYTHON_UNITTEST_MANIFESTS += [
     'taskgraph/test/python.ini',
 ]
 SPHINX_TREES['taskcluster'] = 'docs'
+SPHINX_PYTHON_PACKAGE_DIRS += ['taskgraph']

--- a/taskcluster/taskgraph/morph.py
+++ b/taskcluster/taskgraph/morph.py
@@ -263,22 +263,27 @@ class apply_jsone_templates(object):
         for task in taskgraph.tasks.itervalues():
             for template in sorted(self.templates):
                 context = {
                     'task': task.task,
                     'taskGroup': None,
                     'taskId': task.task_id,
                     'kind': task.kind,
                     'input': self.templates[template],
+                    # The following context differs from action tasks
+                    'attributes': task.attributes,
                 }
 
                 template_path = os.path.join(self.template_dir, template + '.yml')
                 with open(template_path) as f:
                     template = yaml.load(f)
-                task.task = jsone.render(template, context)
+                result = jsone.render(template, context) or {}
+                for attr in ('task', 'attributes'):
+                    if attr in result:
+                        setattr(task, attr, result[attr])
 
         return taskgraph, label_to_taskid
 
 
 def morph(taskgraph, label_to_taskid, parameters):
     """Apply all morphs"""
     morphs = [
         add_index_tasks,

--- a/taskcluster/taskgraph/templates/artifact.yml
+++ b/taskcluster/taskgraph/templates/artifact.yml
@@ -1,18 +1,15 @@
 ---
 $if: task["tags"]
 then:
     $if: task.tags["kind"] == "build"
     then:
-        $mergeDeep:
-            - $eval: task
-            - extra:
-                  treeherder:
-                      symbol: Ba
-            - payload:
-                  env:
-                      USE_ARTIFACT:
-                          $eval: input.enabled
-    else:
-        $eval: task
-else:
-    $eval: task
+        task:
+            $mergeDeep:
+                - $eval: task
+                - extra:
+                      treeherder:
+                          symbol: Ba
+                - payload:
+                      env:
+                          USE_ARTIFACT:
+                              $eval: input.enabled

--- a/taskcluster/taskgraph/templates/env.yml
+++ b/taskcluster/taskgraph/templates/env.yml
@@ -1,5 +1,6 @@
-$mergeDeep:
-    - $eval: task
-    - payload:
-          env:
-              $eval: input
+task:
+    $mergeDeep:
+        - $eval: task
+        - payload:
+              env:
+                  $eval: input