feat: use primary-dependency attribute to find primary dependency of a task

Storing this on `from_deps` tasks allows us to avoid iterating over potentially all the task dependencies when finding it later, which can add up to a lot of time in larger graphs.

Author: bhearsum

src/taskgraph/transforms/from_deps.py

@@ -206,6 +206,7 @@ def from_deps(config, tasks):
             )
 
             primary_dep = [dep for dep in group if dep.kind == primary_kind][0]
+            new_task["attributes"]["primary-dependency-label"] = primary_dep.label
 
             if set_name:
                 func = SET_NAME_MAP[set_name]

src/taskgraph/util/dependencies.py

@@ -68,10 +68,9 @@ def get_dependencies(config: TransformConfig, task: Dict) -> Iterator[Task]:
 
 
 def get_primary_dependency(config: TransformConfig, task: Dict) -> Optional[Task]:
-    """Return the ``Task`` object associated with the primary dependency.
-
-    This uses the task's ``primary-kind-dependency`` attribute to find the primary
-    dependency, or returns ``None`` if the attribute is unset.
+    """Return the ``Task`` object associated with the primary dependency,
+    which is assumed to be available in the ``primary-dependency`` attribute.
+    (Which is always the case for tasks created with ``from_deps``.)
 
     Args:
         config (TransformConfig): The ``TransformConfig`` object associated
@@ -83,10 +82,7 @@ def get_primary_dependency(config: TransformConfig, task: Dict) -> Optional[Task
             primary dependency or ``None``.
     """
     try:
-        primary_kind = task["attributes"]["primary-kind-dependency"]
+        label = task["attributes"]["primary-dependency-label"]
+        return config.kind_dependencies_tasks[label]
     except KeyError:
         return None
-
-    for dep in get_dependencies(config, task):
-        if dep.kind == primary_kind:
-            return dep

test/test_transforms_from_deps.py

@@ -3,65 +3,103 @@
 """
 
 from pprint import pprint
+from typing import Dict
 
 import pytest
 from pytest_taskgraph import make_task
 
+from taskgraph.task import Task
 from taskgraph.transforms import from_deps
 
 
+def get_task_label_by_attributes(
+    tasks: Dict[str, Task], attributes: Dict[str, str]
+) -> str:
+    for task in tasks.values():
+        if all([task.attributes[k] == v for k, v in attributes.items()]):
+            return task.label
+
+    return ""
+
+
 def handle_exception(obj, exc=None):
     if exc:
         assert isinstance(obj, exc)
     elif isinstance(obj, Exception):
         raise obj
 
 
-def assert_no_kind_dependencies(e):
+def assert_no_kind_dependencies(_, e):
     handle_exception(e, exc=Exception)
 
 
-def assert_invalid_only_kinds(e):
+def assert_invalid_only_kinds(_, e):
     handle_exception(e, exc=Exception)
 
 
-def assert_defaults(tasks):
+def assert_defaults(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 2
-    assert tasks[0]["attributes"] == {"primary-kind-dependency": "foo"}
+    assert tasks[0]["attributes"] == {
+        "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(deps, {"kind": "foo"}),
+    }
     assert tasks[0]["dependencies"] == {"foo": "a"}
     assert tasks[0]["name"] == "a"
-    assert tasks[1]["attributes"] == {"primary-kind-dependency": "bar"}
+    assert tasks[1]["attributes"] == {
+        "primary-kind-dependency": "bar",
+        "primary-dependency-label": get_task_label_by_attributes(deps, {"kind": "bar"}),
+    }
     assert tasks[1]["dependencies"] == {"bar": "bar-b"}
     assert tasks[1]["name"] == "b"
 
 
 assert_group_by_single = assert_defaults
 
 
-def assert_group_by_attribute(tasks):
+def assert_group_by_attribute(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 2
     assert tasks[0]["dependencies"] == {"foo": "a"}
-    assert tasks[0]["attributes"] == {"primary-kind-dependency": "foo"}
+    assert tasks[0]["attributes"] == {
+        "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(
+            deps, {"kind": "foo", "build-type": "linux"}
+        ),
+    }
     assert tasks[1]["dependencies"] == {"foo": "b", "bar": "c"}
-    assert tasks[1]["attributes"] == {"primary-kind-dependency": "foo"}
+    assert tasks[1]["attributes"] == {
+        "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(
+            deps, {"kind": "foo", "build-type": "win"}
+        ),
+    }
 
 
-def assert_group_by_attribute_dupe(e):
+def assert_group_by_attribute_dupe(_, e):
     handle_exception(e, exc=Exception)
 
 
-def assert_group_by_attribute_dupe_allowed(tasks):
+def assert_group_by_attribute_dupe_allowed(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 2
     assert tasks[0]["dependencies"] == {"a": "a"}
-    assert tasks[0]["attributes"] == {"primary-kind-dependency": "foo"}
+    assert tasks[0]["attributes"] == {
+        "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(
+            deps, {"kind": "foo", "build-type": "linux"}
+        ),
+    }
     assert tasks[1]["dependencies"] == {"b": "b", "c": "c"}
-    assert tasks[1]["attributes"] == {"primary-kind-dependency": "foo"}
+    assert tasks[1]["attributes"] == {
+        "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(
+            deps, {"kind": "foo", "build-type": "win"}
+        ),
+    }
 
 
-def assert_copy_attributes(tasks):
+def assert_copy_attributes(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
 
@@ -70,48 +108,49 @@ def assert_copy_attributes(tasks):
         "build-type": "win",
         "kind": "foo",
         "primary-kind-dependency": "foo",
+        "primary-dependency-label": get_task_label_by_attributes(deps, {"kind": "foo"}),
     }
 
 
-def assert_group_by_all(tasks):
+def assert_group_by_all(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
     assert tasks[0]["dependencies"] == {"foo": "a", "bar": "bar-b"}
 
 
-def assert_group_by_all_dupe_allowed(tasks):
+def assert_group_by_all_dupe_allowed(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
     assert tasks[0]["dependencies"] == {"a": "a", "b": "b", "c": "c"}
 
 
-def assert_dont_set_name(tasks):
+def assert_dont_set_name(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
     assert tasks[0]["name"] == "a-special-name"
 
 
-def assert_dont_set_name_false(tasks):
+def assert_dont_set_name_false(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
     assert tasks[0]["name"] == "a-special-name"
 
 
-def assert_set_name_strip_kind(tasks):
+def assert_set_name_strip_kind(deps, 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):
+def assert_set_name_retain_kind(deps, 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):
+def assert_group_by_all_with_fetch(deps, tasks):
     handle_exception(tasks)
     assert len(tasks) == 1
     assert tasks[0]["dependencies"] == {
@@ -412,4 +451,4 @@ def test_transforms(
 
     param_id = request.node.callspec.id
     assert_func = globals()[f"assert_{param_id}"]
-    assert_func(result)
+    assert_func(deps, result)

test/test_util_dependencies.py

@@ -22,12 +22,12 @@ def test_get_primary_dependency(make_transform_config):
         "foo": make_task("foo", kind="kind1"),
         "bar": make_task("bar", kind="kind2"),
     }
-    config = make_transform_config(kind_dependencies_tasks=dep_tasks)
     task = {
         "attributes": {},
         "dependencies": {"kind1": "foo", "kind2": "bar"},
     }
+    config = make_transform_config(kind_dependencies_tasks=dep_tasks)
     assert get_primary_dependency(config, task) is None
 
-    task["attributes"]["primary-kind-dependency"] = "kind2"
+    task["attributes"]["primary-dependency-label"] = dep_tasks["bar"].label
     assert get_primary_dependency(config, task) == list(dep_tasks.values())[1]