feat(morph): ability for consumers to register custom morphs

src/taskgraph/morph.py

@@ -33,6 +33,12 @@
 logger = logging.getLogger(__name__)
 MAX_ROUTES = 10
 
+registered_morphs = []
+
+
+def register_morph(func):
+    registered_morphs.append(func)
+
 
 def amend_taskgraph(taskgraph, label_to_taskid, to_add):
     """Add the given tasks to the taskgraph, returning a new taskgraph"""
@@ -156,6 +162,7 @@ def make_index_task(parent_task, taskgraph, label_to_taskid, parameters, graph_c
     return task, taskgraph, label_to_taskid
 
 
+@register_morph
 def add_index_tasks(taskgraph, label_to_taskid, parameters, graph_config):
     """
     The TaskCluster queue only allows 10 routes on a task, but we have tasks
@@ -196,8 +203,9 @@ def _get_morph_url():
     return f"{taskgraph_repo}/raw-file/{taskgraph_rev}/src/taskgraph/morph.py"
 
 
+@register_morph
 def add_code_review_task(taskgraph, label_to_taskid, parameters, graph_config):
-    logger.debug("Morphing: adding index tasks")
+    logger.debug("Morphing: adding code review task")
 
     review_config = parameters.get("code-review")
     if not review_config:
@@ -256,12 +264,7 @@ def add_code_review_task(taskgraph, label_to_taskid, parameters, graph_config):
 
 def morph(taskgraph, label_to_taskid, parameters, graph_config):
     """Apply all morphs"""
-    morphs = [
-        add_index_tasks,
-        add_code_review_task,
-    ]
-
-    for m in morphs:
+    for m in registered_morphs:
         taskgraph, label_to_taskid = m(
             taskgraph, label_to_taskid, parameters, graph_config
         )

test/test_morph.py

@@ -72,3 +72,18 @@ def test_make_index_tasks(make_taskgraph, graph_config):
 
     # check the scope summary
     assert index_task.task["scopes"] == ["index:insert-task:gecko.v2.mozilla-central.*"]
+
+
+def test_register_morph(monkeypatch, make_taskgraph):
+    taskgraph, label_to_taskid = make_taskgraph({})
+    monkeypatch.setattr(morph, "registered_morphs", [])
+
+    @morph.register_morph
+    def fake_morph(taskgraph, label_to_taskid, *args):
+        label_to_taskid.setdefault("count", 0)
+        label_to_taskid["count"] += 1
+        return taskgraph, label_to_taskid
+
+    assert label_to_taskid == {}
+    morph.morph(taskgraph, label_to_taskid, None, None)
+    assert label_to_taskid == {"count": 1}