Fix BranchPythonOperator failure when callable returns None (#54991)

`BranchPythonOperator` now properly handles callables that return None
by skipping all downstream tasks, instead of throwing an execution error.
This restores the expected behavior for users who rely on None returns
to skip branches conditionally.

Fixes #54340

Author: kaxil

providers/standard/src/airflow/providers/standard/operators/branch.py

@@ -37,11 +37,17 @@
 class BranchMixIn(SkipMixin):
     """Utility helper which handles the branching as one-liner."""
 
-    def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
+    def do_branch(
+        self, context: Context, branches_to_execute: str | Iterable[str] | None
+    ) -> str | Iterable[str] | None:
         """Implement the handling of branching including logging."""
         self.log.info("Branch into %s", branches_to_execute)
-        branch_task_ids = self._expand_task_group_roots(context["ti"], branches_to_execute)
-        self.skip_all_except(context["ti"], branch_task_ids)
+        if branches_to_execute is None:
+            # When None is returned, skip all downstream tasks
+            self.skip_all_except(context["ti"], None)
+        else:
+            branch_task_ids = self._expand_task_group_roots(context["ti"], branches_to_execute)
+            self.skip_all_except(context["ti"], branch_task_ids)
         return branches_to_execute
 
     def _expand_task_group_roots(
@@ -86,13 +92,13 @@ class BaseBranchOperator(BaseOperator, BranchMixIn):
 
     inherits_from_skipmixin = True
 
-    def choose_branch(self, context: Context) -> str | Iterable[str]:
+    def choose_branch(self, context: Context) -> str | Iterable[str] | None:
         """
         Abstract method to choose which branch to run.
 
         Subclasses should implement this, running whatever logic is
         necessary to choose a branch and returning a task_id or list of
-        task_ids.
+        task_ids. If None is returned, all downstream tasks will be skipped.
 
         :param context: Context dictionary as passed to execute()
         """

providers/standard/tests/unit/standard/operators/test_python.py

@@ -552,6 +552,33 @@ def f():
         ):
             ti.run()
 
+    def test_none_return_value_should_skip_all_downstream(self):
+        """Test that returning None from callable should skip all downstream tasks."""
+        clear_db_runs()
+        with self.dag_maker(self.dag_id, serialized=True):
+
+            def return_none():
+                return None
+
+            branch_op = self.opcls(task_id=self.task_id, python_callable=return_none, **self.default_kwargs())
+            branch_op >> [self.branch_1, self.branch_2]
+
+        dr = self.dag_maker.create_dagrun()
+        if AIRFLOW_V_3_0_1:
+            from airflow.exceptions import DownstreamTasksSkipped
+
+            with pytest.raises(DownstreamTasksSkipped) as dts:
+                self.dag_maker.run_ti(self.task_id, dr)
+
+            # When None is returned, all downstream tasks should be skipped
+            expected_skipped = {("branch_1", -1), ("branch_2", -1)}
+            assert set(dts.value.tasks) == expected_skipped
+        else:
+            self.dag_maker.run_ti(self.task_id, dr)
+            self.assert_expected_task_states(
+                dr, {self.task_id: State.SUCCESS, "branch_1": State.SKIPPED, "branch_2": State.SKIPPED}
+            )
+
     @pytest.mark.parametrize(
         "choice,expected_states",
         [