feat: allow resolving keys in task context

docs/reference/transforms/task_context.rst

@@ -4,8 +4,8 @@ Task Context Transforms
 =======================
 
 The :mod:`taskgraph.transforms.task_context` transform can be used to
-substitute values into any field in a task with data that is not known
-until ``taskgraph`` runs.
+resolve keys and/or substitute values into any field in a task with data
+that is not known until ``taskgraph`` runs.
 
 This data can be provided in a few ways, as described below.
 
@@ -37,7 +37,10 @@ Then create a ``task-context`` section in your task definition, e.g:
 
    tasks:
      build:
-       description: my description {foo}
+       description:
+        by-foo:
+          FOO: super special description for FOO tasks
+          default: my description of {foo}
        task-context:
          from-parameters:
            foo: foo
@@ -50,9 +53,12 @@ are used:
 
 .. code-block:: yaml
 
-   foo: with some extra
+   foo: FOO
+
+...the description will end up with a value of "super special description for FOO tasks".
 
-...the description will end up with a value of "my description with some extra".
+When ``foo`` is set to any other value (eg: ``bar``) the description will end with the
+``default`` value with the value of ``foo`` substituted in (eg: ``my description of bar``).
 
 When using ``from-parameters`` you may also provide an ordered list of keys to
 look for in the parameters, with the first one found being used. For example,
@@ -154,8 +160,11 @@ Finally, the name of the task is added to the context implicitly. For example:
 This will evaluate the description correctly, even though there are no
 ``task-context`` keys defined on the individual tasks.
 
-Precedence
-----------
+Order of Operations & Precedence
+--------------------------------
+
+Keys will be resolved on ``substitution-fields`` first, then substitution
+will be performed on the resolved value.
 
 If the same key is found in multiple places the order of precedence is as
 follows: ``from-parameters``, ``from-object`` keys, ``from-file`` and finally

src/taskgraph/transforms/task_context.py

@@ -1,14 +1,20 @@
 from typing import Optional, Union
 
 from taskgraph.transforms.base import TransformSequence
-from taskgraph.util.schema import Schema
+from taskgraph.util.schema import Schema, resolve_keyed_by
 from taskgraph.util.templates import deep_get, substitute_task_fields
 from taskgraph.util.yaml import load_yaml
 
 
+class ResolveKeyedByConfigOptions(Schema):
+    # forwarded to `resolve_keyed_by`; see that function for details
+    defer: Optional[list[str]] = None
+    enforce_single_match: Optional[bool] = None
+
+
 class TaskContextConfig(Schema):
     # A list of fields in the task to substitute the provided values
-    # into.
+    # into and/or resolve keys in.
     substitution_fields: list[str]
     # Retrieve task context values from parameters. A single
     # parameter may be provided or a list of parameters in
@@ -22,13 +28,15 @@ class TaskContextConfig(Schema):
     from_file: Optional[str] = None
     # Key/value pairs to be used as task context
     from_object: Optional[object] = None
+    resolve_keyed_by_options: Optional[ResolveKeyedByConfigOptions] = None
 
 
 #: Schema for the task_context transforms
 class TaskContextSchema(Schema, forbid_unknown_fields=False, kw_only=True):
     name: Optional[str] = None
-    # `task-context` can be used to substitute values into any field in a
-    # task with data that is not known until `taskgraph` runs.
+    # `task-context` can be used to substitute values into any field or resolve
+    # keys in a field in a task with data that is not known until `taskgraph`
+    # runs.
     #
     # This data can be provided via `from-parameters` or `from-file`,
     # which can pull in values from parameters and a defined yml file
@@ -46,6 +54,9 @@ class TaskContextSchema(Schema, forbid_unknown_fields=False, kw_only=True):
     #   - File
     #
     # That is to say: parameters will always override anything else.
+    #
+    # Each entry in `substitution-fields` will first have `resolved_keyed_by`
+    # called on it, and then substitutions will be performed.
     task_context: Optional[TaskContextConfig] = None
 
 
@@ -89,6 +100,12 @@ def render_task(config, tasks):
         if "name" in task:
             subs.setdefault("name", task["name"])
 
-        # Now that we have our combined context, we can substitute.
+        # Now that we have our combined context, we can resolve keys and substitute.
+        task_name = task.get("name", task.get("label", "?"))
+        resolve_keyed_by_options = {}
+        if opts := sub_config.get("resolve-keyed-by-options"):
+            resolve_keyed_by_options.update(opts)
+        for field in fields:
+            resolve_keyed_by(task, field, task_name, **resolve_keyed_by_options, **subs)
         substitute_task_fields(task, fields, **subs)
         yield task

test/test_transform_task_context.py

@@ -1,5 +1,5 @@
 """
-Tests for the 'fetch' transforms.
+Tests for the 'task_context' transforms.
 """
 
 import os.path
@@ -14,6 +14,38 @@
 
 here = os.path.abspath(os.path.dirname(__file__))
 
+BASE_PARAMS = {
+    "base_repository": "http://hg.example.com",
+    "build_date": 0,
+    "build_number": 1,
+    "enable_always_target": True,
+    "head_repository": "http://hg.example.com",
+    "head_rev": "abcdef",
+    "head_ref": "default",
+    "level": "1",
+    "moz_build_date": 0,
+    "next_version": "1.0.1",
+    "owner": "some-owner",
+    "project": "some-project",
+    "pushlog_id": 1,
+    "repository_type": "hg",
+    "target_tasks_method": "test_method",
+    "tasks_for": "hg-push",
+    "try_mode": None,
+    "version": "1.0.0",
+}
+
+
+def make_config(graph_config, extra_params=None):
+    params_data = dict(BASE_PARAMS)
+    if extra_params:
+        params_data.update(extra_params)
+    params = FakeParameters(params_data)
+    return TransformConfig(
+        "test", str(here), {}, params, {}, graph_config, write_artifacts=False
+    )
+
+
 TASK_DEFAULTS = {
     "description": "fake description {object} {file} {param} {object_and_file} "
     "{object_and_param} {file_and_param} {object_file_and_param} {param_fallback} {name}",
@@ -56,39 +88,135 @@
     ),
 )
 def test_transforms(request, run_transform, graph_config, task, description):
-    params = FakeParameters(
+    config = make_config(
+        graph_config,
         {
             "param": "param",
             "object_and_param": "param-overrides-object",
             "file_and_param": "param-overrides-file",
             "object_file_and_param": "param-overrides-all",
             "default": "default",
-            "base_repository": "http://hg.example.com",
-            "build_date": 0,
-            "build_number": 1,
-            "enable_always_target": True,
-            "head_repository": "http://hg.example.com",
-            "head_rev": "abcdef",
-            "head_ref": "default",
-            "level": "1",
-            "moz_build_date": 0,
-            "next_version": "1.0.1",
-            "owner": "some-owner",
-            "project": "some-project",
-            "pushlog_id": 1,
-            "repository_type": "hg",
-            "target_tasks_method": "test_method",
-            "tasks_for": "hg-push",
-            "try_mode": None,
-            "version": "1.0.0",
         },
     )
-    transform_config = TransformConfig(
-        "test", str(here), {}, params, {}, graph_config, write_artifacts=False
-    )
 
-    task = run_transform(task_context.transforms, task, config=transform_config)[0]
+    task = run_transform(task_context.transforms, task, config=config)[0]
     print("Dumping task:")
     pprint(task, indent=2)
 
     assert task["description"] == description
+
+
+def test_resolve_keyed_by_from_parameters(run_transform, graph_config):
+    task = {
+        "name": "fake-task-name",
+        "description": {
+            "by-platform": {
+                "linux": "resolved-from-parameters",
+                "default": "default-value",
+            },
+        },
+        "task-context": {
+            "from-parameters": {"platform": "platform"},
+            "substitution-fields": ["description"],
+        },
+    }
+    config = make_config(graph_config, {"platform": "linux"})
+
+    task = run_transform(task_context.transforms, task, config=config)[0]
+    pprint(task, indent=2)
+
+    assert task["description"] == "resolved-from-parameters"
+
+
+def test_resolve_keyed_by_from_file(run_transform, graph_config):
+    task = {
+        "name": "fake-task-name",
+        "description": {
+            "by-file": {
+                "file": "resolved-from-file",
+                "default": "default-value",
+            },
+        },
+        "task-context": {
+            "from-file": f"{here}/data/task_context.yml",
+            "substitution-fields": ["description"],
+        },
+    }
+    config = make_config(graph_config)
+
+    task = run_transform(task_context.transforms, task, config=config)[0]
+    pprint(task, indent=2)
+
+    assert task["description"] == "resolved-from-file"
+
+
+def test_resolve_keyed_by_from_object(run_transform, graph_config):
+    task = {
+        "name": "fake-task-name",
+        "description": {
+            "by-flavor": {
+                "chocolate": "resolved-from-object",
+                "default": "default-value",
+            },
+        },
+        "task-context": {
+            "from-object": {"flavor": "chocolate"},
+            "substitution-fields": ["description"],
+        },
+    }
+    config = make_config(graph_config)
+
+    task = run_transform(task_context.transforms, task, config=config)[0]
+    pprint(task, indent=2)
+
+    assert task["description"] == "resolved-from-object"
+
+
+def test_resolve_keyed_by_missing_context_no_defer(run_transform, graph_config):
+    task = {
+        "name": "fake-task-name",
+        "description": {
+            "by-missing": {
+                "linux": "should-not-resolve",
+            },
+        },
+        "task-context": {
+            "substitution-fields": ["description"],
+        },
+    }
+    config = make_config(graph_config)
+
+    with pytest.raises(Exception):
+        run_transform(task_context.transforms, task, config=config)
+
+
+def test_resolve_keyed_by_missing_context_with_defer(run_transform, graph_config):
+    task = {
+        "name": "fake-task-name",
+        "description": {
+            "by-flavor": {
+                "chocolate": {
+                    "by-missing": {
+                        "linux": "should-not-resolve",
+                        "default": "also-should-not-resolve",
+                    },
+                },
+            },
+        },
+        "task-context": {
+            "from-object": {"flavor": "chocolate"},
+            "resolve-keyed-by-options": {"defer": ["missing"]},
+            "substitution-fields": ["description"],
+        },
+    }
+    config = make_config(graph_config)
+
+    task = run_transform(task_context.transforms, task, config=config)[0]
+    pprint(task, indent=2)
+
+    assert task["description"] == {
+        "by-missing": {
+            "linux": "should-not-resolve",
+            "default": "also-should-not-resolve",
+        },
+    }