Add a function to get tasks from multiple kinds in one pass

packages/pytest-taskgraph/src/pytest_taskgraph/fixtures/gen.py

@@ -30,7 +30,7 @@ def fake_loader(kind, path, config, parameters, loaded_tasks):
             "attributes": {"_tasknum": str(i)},
             "task": {
                 "i": i,
-                "metadata": {"name": f"t-{i}"},
+                "metadata": {"name": f"{kind}-t-{i}"},
                 "deadline": "soon",
             },
             "dependencies": dependencies,

src/taskgraph/generator.py

@@ -446,19 +446,28 @@ def verify(self, name, *args, **kwargs):
             return name, args[0]
 
 
-def load_tasks_for_kind(parameters, kind, root_dir=None):
+def load_tasks_for_kinds(parameters, kinds, root_dir=None):
     """
-    Get all the tasks of a given kind.
+    Get all the tasks of the given kinds.
 
     This function is designed to be called from outside of taskgraph.
     """
     # make parameters read-write
     parameters = dict(parameters)
-    parameters["target-kinds"] = [kind]
+    parameters["target-kinds"] = kinds
     parameters = parameters_loader(spec=None, strict=False, overrides=parameters)
     tgg = TaskGraphGenerator(root_dir=root_dir, parameters=parameters)
     return {
         task.task["metadata"]["name"]: task
         for task in tgg.full_task_set
-        if task.kind == kind
+        if task.kind in kinds
     }
+
+
+def load_tasks_for_kind(parameters, kind, root_dir=None):
+    """
+    Get all the tasks of a given kind.
+
+    This function is designed to be called from outside of taskgraph.
+    """
+    return load_tasks_for_kinds(parameters, [kind], root_dir)

test/test_generator.py

@@ -7,7 +7,7 @@
 from pytest_taskgraph import FakeKind, WithFakeKind, fake_load_graph_config
 
 from taskgraph import generator, graph
-from taskgraph.generator import Kind, load_tasks_for_kind
+from taskgraph.generator import Kind, load_tasks_for_kind, load_tasks_for_kinds
 from taskgraph.loader.default import loader as default_loader
 
 
@@ -184,7 +184,25 @@ def test_load_tasks_for_kind(monkeypatch):
         "_example-kind",
         "/root",
     )
-    assert "t-1" in tasks and tasks["t-1"].label == "_example-kind-t-1"
+    assert "docker-image-t-1" not in tasks
+    assert (
+        "_example-kind-t-1" in tasks
+        and tasks["_example-kind-t-1"].label == "_example-kind-t-1"
+    )
+
+    tasks = load_tasks_for_kinds(
+        {"_kinds": [("_example-kind", []), ("docker-image", [])]},
+        ["_example-kind", "docker-image"],
+        "/root",
+    )
+    assert (
+        "docker-image-t-1" in tasks
+        and tasks["docker-image-t-0"].label == "docker-image-t-0"
+    )
+    assert (
+        "_example-kind-t-1" in tasks
+        and tasks["_example-kind-t-1"].label == "_example-kind-t-1"
+    )
 
 
 @pytest.mark.parametrize(