This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new aa54031f3ce Allow accessing a TaskGroup's members via `[]` (#64430)
aa54031f3ce is described below
commit aa54031f3cece54b1492bb13648bd521d5bf3082
Author: Dev-iL <[email protected]>
AuthorDate: Tue Apr 21 09:02:24 2026 +0300
Allow accessing a TaskGroup's members via `[]` (#64430)
---
task-sdk/src/airflow/sdk/definitions/taskgroup.py | 4 ++++
.../tests/task_sdk/definitions/test_taskgroup.py | 25 ++++++++++++++++++++++
2 files changed, 29 insertions(+)
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index c47b1f360ae..72779adb5e8 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -141,6 +141,10 @@ class TaskGroup(DAGNode):
if not dag:
raise RuntimeError("TaskGroup can only be used inside a dag")
+ def __getitem__(self, key: str) -> DAGNode:
+ """Return a child node by label via ``tg[label]`` syntax. See
`get_child_by_label`."""
+ return self.get_child_by_label(key)
+
def __attrs_post_init__(self):
# TODO: If attrs supported init only args we could use that here
# https://github.com/python-attrs/attrs/issues/342
diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
index 6eb51a1a2f6..5eb69dddd4c 100644
--- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
+++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
@@ -223,6 +223,31 @@ def test_build_task_group_with_prefix():
assert group4.get_child_by_label("task4") == task4
[email protected](
+ "prefix",
+ [
+ pytest.param(True, id="prefix_on"),
+ pytest.param(False, id="prefix_off"),
+ ],
+)
+def test_taskgroup_getitem_returns_child_by_label(prefix: bool):
+ """Tests that TaskGroup[label] returns the correct child task or
subgroup."""
+ logical_date = pendulum.datetime(2020, 1, 1)
+ with DAG("test_getitem", start_date=logical_date):
+ with TaskGroup("group1", prefix_group_id=prefix) as group1:
+ task1 = EmptyOperator(task_id="task1")
+ with TaskGroup("subgroup", prefix_group_id=prefix) as subgroup:
+ task2 = EmptyOperator(task_id="task2")
+
+ assert group1["task1"] == task1
+ assert group1["subgroup"] == subgroup
+ assert group1["subgroup"]["task2"] == task2
+
+ # Missing label raises KeyError
+ with pytest.raises(KeyError):
+ group1["nonexistent"]
+
+
def test_build_task_group_with_prefix_functionality():
"""
Tests TaskGroup prefix_group_id functionality - additional test for
comprehensive coverage.
