uranusjr commented on code in PR #53696:
URL: https://github.com/apache/airflow/pull/53696#discussion_r2244766058
##########
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:
This needs to match the function name in sdk. We can change it (IIRC we
didn’t use a property because there used to be a function called
`needs_expansion`), but I’d prefer this to be done when we move the logic into
_shared instead to minimise unrelated diff.
--
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]
