ephraimbuddy commented on code in PR #32235:
URL: https://github.com/apache/airflow/pull/32235#discussion_r1253335566
##########
airflow/utils/setup_teardown.py:
##########
@@ -63,29 +65,37 @@ def pop_context_managed_setup_task(cls) -> Operator |
list[Operator] | None:
else:
setup_task.set_downstream(old_setup_task)
else:
- cls._context_managed_setup_task = None
+ cls._context_managed_setup_task = []
return old_setup_task
@classmethod
- def update_context_map(cls, operator):
- ctx = BaseSetupTeardownContext.context_map
- if setup_task :=
BaseSetupTeardownContext.get_context_managed_setup_task():
- if isinstance(setup_task, list):
- setup_task = tuple(setup_task)
- if ctx.get(setup_task) is None:
- ctx[setup_task] = [operator]
+ def update_context_map(cls, task: DependencyMixin):
+ from airflow.models.abstractoperator import AbstractOperator
+
+ task_ = cast(AbstractOperator, task)
+ if task_.is_setup or task_.is_teardown:
+ return
+ ctx = cls.context_map
+
+ def _get_or_set_item(item):
+ if ctx.get(item) is None:
+ ctx[item] = [task_]
else:
- ctx[setup_task].append(operator)
- if teardown_task :=
BaseSetupTeardownContext.get_context_managed_teardown_task():
+ ctx[item].append(task_)
+
+ if setup_task := cls.get_context_managed_setup_task():
+ if isinstance(setup_task, list):
+ _get_or_set_item(tuple(setup_task))
+ elif not isinstance(setup_task, list):
Review Comment:
I think I had issues with mypy but I will check again
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
