Lee-W commented on code in PR #53696:
URL: https://github.com/apache/airflow/pull/53696#discussion_r2238776303
##########
airflow-core/src/airflow/models/mappedoperator.py:
##########
@@ -152,6 +341,104 @@ def get_extra_links(self, ti: TaskInstance, name: str) ->
str | None:
return None
return link.get_link(self, ti_key=ti.key)
+ # TODO (GH-52141): Copied from sdk. Find a better place for this to live
in.
+ def _get_specified_expand_input(self) -> SchedulerExpandInput:
+ """Input received from the expand call on the operator."""
+ return getattr(self, self._expand_input_attr)
+
+ # TODO (GH-52141): Copied from sdk. Find a better place for this to live
in.
+ def iter_mapped_task_groups(self) -> Iterator[MappedTaskGroup]:
+ """
+ Return mapped task groups this task belongs to.
+
+ Groups are returned from the innermost to the outmost.
+
+ :meta private:
+ """
+ if (group := self.task_group) is None:
+ return
+ yield from group.iter_mapped_task_groups()
+
+ # TODO (GH-52141): Copied from sdk. Find a better place for this to live
in.
+ def get_closest_mapped_task_group(self) -> MappedTaskGroup | None:
+ """
+ Get the mapped task group "closest" to this task in the DAG.
+
+ :meta private:
+ """
+ return next(self.iter_mapped_task_groups(), None)
+
+ # TODO (GH-52141): Copied from sdk. Find a better place for this to live
in.
+ def get_needs_expansion(self) -> bool:
Review Comment:
Would it be better to make it a property? probably can be done in 52141
##########
airflow-core/src/airflow/models/taskmap.py:
##########
@@ -131,14 +131,14 @@ def expand_mapped_task(cls, task, run_id: str, *,
session: Session) -> tuple[Seq
order by map index, and the maximum map index value.
"""
from airflow.models.expandinput import NotFullyPopulated
- from airflow.models.mappedoperator import get_mapped_ti_count
+ from airflow.models.mappedoperator import MappedOperator as
DbMappedOperator, get_mapped_ti_count
Review Comment:
Should we unify the way we name the ones from `sdk.definitions`,
`sdk.types`, or `models` and document it? Probably don't need to be done in
this PR, but it could be a bit confusing when reading through different
modules.
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -147,10 +147,8 @@ def _get_expanded_ti_count() -> int:
return get_mapped_ti_count(ti.task, ti.run_id, session=session)
def _iter_expansion_dependencies(task_group: MappedTaskGroup) ->
Iterator[str]:
- from airflow.sdk.definitions.mappedoperator import MappedOperator
-
- if isinstance(ti.task, MappedOperator):
- for op in ti.task.iter_mapped_dependencies():
+ if (task := ti.task) is not None and is_mapped(task):
+ for op in task.iter_mapped_dependencies():
yield op.task_id
Review Comment:
```suggestion
yield from (op.task_id for op in
task.iter_mapped_dependencies())
```
--
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]
