Serban Stanca
9e40b6bb18
This reverts commit540cd75043. This reverts commitafef97fa49. This reverts commit3f4efe4bbd.
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
# 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 collections.abc import Iterator
|
|
from typing import 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,
|
|
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
|
|
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:
|
|
label = task["attributes"]["primary-dependency-label"]
|
|
return config.kind_dependencies_tasks[label]
|
|
except KeyError:
|
|
return None
|