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]