Source code
Revision control
Copy as Markdown
Other Tools
# 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
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