def transform(self, input_operator: BaseOperator,
parent_fragment: DAGFragment,
upstream_fragments: List[DAGFragment]) -> DAGFragment:
input_operator._dag = None
input_operator._upstream_task_ids.clear()
input_operator._downstream_task_ids.clear()
input_operator.dag = self.dag
return DAGFragment([input_operator])
def remove_task_from_dag(dag: DAG, dag_nodes: List[BaseOperator],
task: BaseOperator):
"""
Removes the given list of :class:`~airflow.models.BaseOperator`\'s from the given
:class:`~airflow.models.DAG`
:param dag: the source airflow DAG
:param dag_nodes: the list of nodes in the source DAG
:param task: the task to remove
"""
all_other_tasks = [t for t in dag_nodes if t is not task]
for this_task in all_other_tasks:
if task.task_id in this_task._upstream_task_ids:
this_task._upstream_task_ids.remove(task.task_id)
if task.task_id in this_task._downstream_task_ids:
this_task._downstream_task_ids.remove(task.task_id)
task._upstream_task_ids.clear()
task._downstream_task_ids.clear()
task._dag = None
del dag.task_dict[task.task_id]