summaryrefslogtreecommitdiffstats
path: root/third_party/python/taskcluster_taskgraph/taskgraph/util/dependencies.py
blob: d33aa3d7f28f0203eda5cf62aaf003027d963a66 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
# 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/.

from typing import Dict, Iterator, Optional

from taskgraph.task import Task
from taskgraph.transforms.base import TransformConfig
from taskgraph.util.schema import Schema

# Define a collection of group_by functions
GROUP_BY_MAP = {}


def group_by(name, schema=None):
    def wrapper(func):
        assert (
            name not in GROUP_BY_MAP
        ), f"duplicate group_by function name {name} ({func} and {GROUP_BY_MAP[name]})"
        GROUP_BY_MAP[name] = func
        func.schema = schema
        return func

    return wrapper


@group_by("single")
def group_by_single(config, tasks):
    for task in tasks:
        yield [task]


@group_by("all")
def group_by_all(config, tasks):
    return [[task for task in tasks]]


@group_by("attribute", schema=Schema(str))
def group_by_attribute(config, tasks, attr):
    groups = {}
    for task in tasks:
        val = task.attributes.get(attr)
        if not val:
            continue
        groups.setdefault(val, []).append(task)

    return groups.values()


def get_dependencies(config: TransformConfig, task: Dict) -> Iterator[Task]:
    """Iterate over all dependencies as ``Task`` objects.

    Args:
        config (TransformConfig): The ``TransformConfig`` object associated
            with the kind.
        task (Dict): The task dictionary to retrieve dependencies from.

    Returns:
        Iterator[Task]: Returns a generator that iterates over the ``Task``
        objects associated with each dependency.
    """
    if "dependencies" not in task:
        return []

    for label, dep in config.kind_dependencies_tasks.items():
        if label in task["dependencies"].values():
            yield dep


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.

    Args:
        config (TransformConfig): The ``TransformConfig`` object associated
            with the kind.
        task (Dict): The task dictionary to retrieve the primary dependency from.

    Returns:
        Optional[Task]: The ``Task`` object associated with the
            primary dependency or ``None``.
    """
    try:
        primary_kind = task["attributes"]["primary-kind-dependency"]
    except KeyError:
        return None

    for dep in get_dependencies(config, task):
        if dep.kind == primary_kind:
            return dep