feat: more featureful set-name for from-deps (#484)

* Replace `set-name` with a more featureful version

* style: pre-commit.ci auto fixes [...]

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bhearsum

src/taskgraph/transforms/from_deps.py

@@ -20,6 +20,7 @@
 from taskgraph.util.attributes import attrmatch
 from taskgraph.util.dependencies import GROUP_BY_MAP, get_dependencies
 from taskgraph.util.schema import Schema, validate_schema
+from taskgraph.util.set_name import SET_NAME_MAP
 
 FROM_DEPS_SCHEMA = Schema(
     {
@@ -41,12 +42,14 @@
                 "set-name",
                 description=dedent(
                     """
-                When True, `from_deps` will derive a name for the generated
-                tasks from the name of the primary dependency. Defaults to
-                True.
+                UPDATE ME AND DOCS
                 """.lstrip()
                 ),
-            ): bool,
+            ): Any(
+                None,
+                *SET_NAME_MAP,
+                {Any(*SET_NAME_MAP): object},
+            ),
             Optional(
                 "with-attributes",
                 description=dedent(
@@ -170,7 +173,7 @@ def from_deps(config, tasks):
             groups = func(config, deps)
 
         # Split the task, one per group.
-        set_name = from_deps.get("set-name", True)
+        set_name = from_deps.get("set-name", "strip-kind")
         copy_attributes = from_deps.get("copy-attributes", False)
         unique_kinds = from_deps.get("unique-kinds", True)
         fetches = from_deps.get("fetches", [])
@@ -203,10 +206,8 @@ def from_deps(config, tasks):
             primary_dep = [dep for dep in group if dep.kind == primary_kind][0]
 
             if set_name:
-                if primary_dep.label.startswith(primary_kind):
-                    new_task["name"] = primary_dep.label[len(primary_kind) + 1 :]
-                else:
-                    new_task["name"] = primary_dep.label
+                func = SET_NAME_MAP[set_name]
+                new_task["name"] = func(config, deps, primary_dep, primary_kind)
 
             if copy_attributes:
                 attrs = new_task.setdefault("attributes", {})

src/taskgraph/util/set_name.py

@@ -0,0 +1,34 @@
+# 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 set_name functions
+# Note: this is stored here instead of where it is used in the `from_deps`
+# transform to give consumers a chance to register their own `set_name`
+# handlers before the `from_deps` schema is created.
+SET_NAME_MAP = {}
+
+
+def set_name(name, schema=None):
+    def wrapper(func):
+        assert (
+            name not in SET_NAME_MAP
+        ), f"duplicate set_name function name {name} ({func} and {SET_NAME_MAP[name]})"
+        SET_NAME_MAP[name] = func
+        func.schema = schema
+        return func
+
+    return wrapper
+
+
+@set_name("strip-kind")
+def set_name_strip_kind(config, tasks, primary_dep, primary_kind):
+    if primary_dep.label.startswith(primary_kind):
+        return primary_dep.label[len(primary_kind) + 1 :]
+    else:
+        return primary_dep.label
+
+
+@set_name("retain-kind")
+def set_name_retain_kind(config, tasks, primary_dep, primary_kind):
+    return primary_dep.label

test/test_transforms_from_deps.py

@@ -92,6 +92,20 @@ def assert_dont_set_name(tasks):
     assert tasks[0]["name"] == "a-special-name"
 
 
+def assert_set_name_strip_kind(tasks):
+    handle_exception(tasks)
+    assert len(tasks) == 2
+    assert tasks[0]["name"] == "a"
+    assert tasks[1]["name"] == "b"
+
+
+def assert_set_name_retain_kind(tasks):
+    handle_exception(tasks)
+    assert len(tasks) == 2
+    assert tasks[0]["name"] == "a"
+    assert tasks[1]["name"] == "bar-b"
+
+
 def assert_group_by_all_with_fetch(tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
@@ -165,7 +179,7 @@ def assert_group_by_all_with_fetch(tasks):
                 "name": "a-special-name",
                 "from-deps": {
                     "group-by": "all",
-                    "set-name": False,
+                    "set-name": None,
                 },
             },
             # kind config
@@ -174,6 +188,32 @@ def assert_group_by_all_with_fetch(tasks):
             None,
             id="dont_set_name",
         ),
+        pytest.param(
+            # task
+            {
+                "from-deps": {
+                    "set-name": "strip-kind",
+                },
+            },
+            # kind config
+            None,
+            # deps
+            None,
+            id="set_name_strip_kind",
+        ),
+        pytest.param(
+            # task
+            {
+                "from-deps": {
+                    "set-name": "retain-kind",
+                },
+            },
+            # kind config
+            None,
+            # deps
+            None,
+            id="set_name_retain_kind",
+        ),
         pytest.param(
             # task
             {"from-deps": {}},