[airflow] branch main updated: Add override method to TaskGroupDecorator (#25160)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk 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 b815e221e8 Add override method to TaskGroupDecorator (#25160)
b815e221e8 is described below

commit b815e221e828aa039e166378c3e393d144b174e8
Author: Denis Makarov <44...@users.noreply.github.com>
AuthorDate: Sat Jul 23 00:30:27 2022 +0300

    Add override method to TaskGroupDecorator (#25160)
---
 airflow/decorators/task_group.py    |  3 +++
 tests/decorators/test_task_group.py | 52 +++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/airflow/decorators/task_group.py b/airflow/decorators/task_group.py
index 6744749470..59adad95ea 100644
--- a/airflow/decorators/task_group.py
+++ b/airflow/decorators/task_group.py
@@ -80,6 +80,9 @@ class TaskGroupDecorator(Generic[R]):
         #   start >> tg >> end
         return task_group
 
+    def override(self, **kwargs: Any) -> "TaskGroupDecorator[R]":
+        return attr.evolve(self, kwargs={**self.kwargs, **kwargs})
+
 
 class Group(Generic[F]):
     """Declaration of a @task_group-decorated callable for type-checking.
diff --git a/tests/decorators/test_task_group.py b/tests/decorators/test_task_group.py
new file mode 100644
index 0000000000..2d568749dc
--- /dev/null
+++ b/tests/decorators/test_task_group.py
@@ -0,0 +1,52 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.decorators import task_group
+
+
+def test_task_group_with_overridden_kwargs():
+    @task_group(
+        default_args={
+            'params': {
+                'x': 5,
+                'y': 5,
+            },
+        },
+        add_suffix_on_collision=True,
+    )
+    def simple_tg():
+        ...
+
+    tg_with_overridden_kwargs = simple_tg.override(
+        group_id='custom_group_id',
+        default_args={
+            'params': {
+                'x': 10,
+            },
+        },
+    )
+
+    assert tg_with_overridden_kwargs.kwargs == {
+        'group_id': 'custom_group_id',
+        'default_args': {
+            'params': {
+                'x': 10,
+            },
+        },
+        'add_suffix_on_collision': True,
+    }