feat: add support for custom context to task-context

docs/reference/transforms/task_context.rst

@@ -98,35 +98,6 @@ with the this ``kind``:
 ...the description will bring in the value ``foo`` from the parameters if
 present, or ``default`` otherwise.
 
-from-file
-~~~~~~~~~
-
-Context may also be provided from a defined yaml file. The provided file
-should usually only contain top level keys and values (eg: nested objects
-will not be interpolated - they will be substituted as text representations
-of the object).
-
-For example, with this kind definition:
-
-.. code-block:: yaml
-
-   tasks:
-     build:
-       description: my description {foo}
-       task-context:
-         from-file: some_file.yaml
-         substitution-fields:
-           - description
-
-And this in ``some_file.yaml``:
-
-.. code-block:: yaml
-
-   foo: from a file
-
-...description will end up with "my description from a file".
-
-
 from-object
 ~~~~~~~~~~~
 
@@ -158,6 +129,86 @@ For example:
 This will give build1 and build2 descriptions with their ``extra_desc``
 included while allowing them to share the rest of their task definition.
 
+from-custom
+~~~~~~~~~~~
+
+Context may be provided by custom providers, which must be registered prior
+to this transform being run. This allows the creation of context that is
+derived from parameters, code constants, or anything else accessible to
+taskgraph.
+
+For example, you may have a custom context handler set-up in
+``taskcluster/my_taskgraph/custom_context.py``:
+
+.. code-block:: python
+
+   NON_PRODUCTION_BRANCHES = ["maple", "pine"]
+
+   @custom_context("release-level")
+   def release_level_context(config, task):
+       # Despite being level 3, some branches are not truly considered "production"
+       # in the sense of creating releases that ship to users.
+       if config.params["level"] == "1" or config.params["project"] in NON_PRODUCTION_BRANCHES:
+           return "staging"
+       return "production"
+
+
+In your ``register`` function you will need to ensure you import this module.
+Eg: ``taskcluster/my_taskgraph/__init__.py``:
+
+.. code-block:: python
+
+   def register(graph_config):
+       from my_taskgraph import custom_contexts  # trigger custom task-context registration
+
+
+Now you can use ``release-level`` as context in a kind:
+
+.. code-block:: yaml
+
+   task-defaults:
+     task-context:
+       from-custom:
+         - release-level
+       substitution-fields:
+         - scopes
+
+     scopes:
+       by-release-level:
+         staging:
+           - secrets:get:staging_creds
+         production:
+           - secrets:get:production_creds
+
+
+from-file
+~~~~~~~~~
+
+Context may also be provided from a defined yaml file. The provided file
+should usually only contain top level keys and values (eg: nested objects
+will not be interpolated - they will be substituted as text representations
+of the object).
+
+For example, with this kind definition:
+
+.. code-block:: yaml
+
+   tasks:
+     build:
+       description: my description {foo}
+       task-context:
+         from-file: some_file.yaml
+         substitution-fields:
+           - description
+
+And this in ``some_file.yaml``:
+
+.. code-block:: yaml
+
+   foo: from a file
+
+...description will end up with "my description from a file".
+
 Implicit Context
 ~~~~~~~~~~~~~~~~
 
@@ -185,7 +236,7 @@ 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
+follows: ``from-parameters``, ``from-object`` keys, ``from-custom`` providers, ``from-file`` and finally
 implicit context.
 
 That is to say: parameters will always override anything else.

src/taskgraph/transforms/task_context.py

@@ -2,6 +2,7 @@
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.schema import Schema, resolve_keyed_by
+from taskgraph.util.task_context import CUSTOM_CONTEXT_MAP
 from taskgraph.util.templates import deep_get, substitute_task_fields
 from taskgraph.util.yaml import load_yaml
 
@@ -28,6 +29,12 @@ class TaskContextConfig(Schema):
     from_file: Optional[str] = None
     # Key/value pairs to be used as task context
     from_object: Optional[object] = None
+    # Retrieve task context values from registered custom providers. Each
+    # entry is the name of a provider registered via
+    # ``taskgraph.util.task_context.custom_context``. Providers are called
+    # with the ``TransformConfig`` and the task being rendered, and must
+    # return a dict of key/value pairs to add to the context.
+    from_custom: Optional[list[str]] = None
     resolve_keyed_by_options: Optional[ResolveKeyedByConfigOptions] = None
 
 
@@ -51,6 +58,7 @@ class TaskContextSchema(Schema, forbid_unknown_fields=False, kw_only=True):
     # is as follows:
     #   - Parameters
     #   - `from-object` keys
+    #   - Custom providers (`from-custom`)
     #   - File
     #
     # That is to say: parameters will always override anything else.
@@ -89,10 +97,15 @@ def render_task(config, tasks):
         if from_file:
             file_context = load_yaml(from_file)
 
+        custom_context = {}
+        for name in sub_config.pop("from-custom", None) or []:
+            custom_context.update(CUSTOM_CONTEXT_MAP[name](config, task))
+
         fields = sub_config.pop("substitution-fields")
 
         subs = {}
         subs.update(file_context)
+        subs.update(custom_context)
         # We've popped away the configuration; everything left in `sub_config` is
         # substitution key/value pairs.
         subs.update(sub_config.pop("from-object", {}))

src/taskgraph/util/task_context.py

@@ -0,0 +1,20 @@
+# 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/.
+
+# Define a collection of custom task-context providers.
+# Note: this is stored here instead of where it is used in the
+# `task_context` transform to give consumers a chance to register their own
+# providers before the `task_context` schema is created.
+CUSTOM_CONTEXT_MAP = {}
+
+
+def custom_context(name):
+    def wrapper(func):
+        assert name not in CUSTOM_CONTEXT_MAP, (
+            f"duplicate custom_context function name {name} ({func} and {CUSTOM_CONTEXT_MAP[name]})"
+        )
+        CUSTOM_CONTEXT_MAP[name] = func
+        return func
+
+    return wrapper

test/test_transform_task_context.py

@@ -11,6 +11,7 @@
 
 from taskgraph.transforms import task_context
 from taskgraph.transforms.base import TransformConfig
+from taskgraph.util.task_context import CUSTOM_CONTEXT_MAP, custom_context
 
 here = os.path.abspath(os.path.dirname(__file__))
 
@@ -220,3 +221,35 @@ def test_resolve_keyed_by_missing_context_with_defer(run_transform, graph_config
             "default": "also-should-not-resolve",
         },
     }
+
+
+def test_resolve_keyed_by_custom_context(run_transform, graph_config):
+    try:
+
+        @custom_context("colour")
+        def _colour(config, task):
+            return {"colour": "blue"}
+
+        task = {
+            "name": "fake-task-name",
+            "description": {
+                "by-colour": {
+                    "blue": "i am a blue task!",
+                    "default": " i am a default task",
+                },
+            },
+            "task-context": {
+                "from-custom": ["colour"],
+                "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"] == "i am a blue task!"
+    finally:
+        # ensure `CUSTOM_CONTEXT_MAP` is reset at the end of the test
+        # regardless of result
+        CUSTOM_CONTEXT_MAP.pop("colour", None)