uranusjr commented on a change in pull request #20411:
URL: https://github.com/apache/airflow/pull/20411#discussion_r773562188
##########
File path: airflow/operators/python.py
##########
@@ -217,8 +217,10 @@ def execute(self, context: Context) -> Any:
branches = {branch}
elif isinstance(branch, list):
branches = set(branch)
+ elif branch is None:
+ branches = set()
else:
- raise AirflowException("Branch callable must return either a task
ID or a list of IDs")
+ raise AirflowException("Branch callable must return either None or
a task ID or a list of IDs")
Review comment:
```suggestion
raise AirflowException("Branch callable must return either None,
a task ID, or a list of IDs")
```
##########
File path: airflow/models/skipmixin.py
##########
@@ -140,6 +140,8 @@ def skip_all_except(self, ti: TaskInstance,
branch_task_ids: Union[str, Iterable
self.log.info("Following branch %s", branch_task_ids)
if isinstance(branch_task_ids, str):
branch_task_ids = {branch_task_ids}
+ elif branch_task_ids is None:
+ branch_task_ids = {}
Review comment:
```suggestion
branch_task_ids = ()
```
Creating a dict is needless.
##########
File path: airflow/models/skipmixin.py
##########
@@ -128,7 +128,7 @@ def skip(
session=session,
)
- def skip_all_except(self, ti: TaskInstance, branch_task_ids: Union[str,
Iterable[str]]):
+ def skip_all_except(self, ti: TaskInstance, branch_task_ids:
Optional[Union[str, Iterable[str]]]):
Review comment:
```suggestion
def skip_all_except(self, ti: TaskInstance, branch_task_ids: Union[None,
str, Iterable[str]]):
```
##########
File path: tests/models/test_skipmixin.py
##########
@@ -117,3 +117,51 @@ def get_state(ti):
assert get_state(ti2) == State.NONE
assert get_state(ti3) == State.SKIPPED
+
+ def test_skip_all_except_none(self, dag_maker):
Review comment:
You can avoid all these new tests (which duplicate most code in
`test_skip_all_except`) with `pytest.mark.parametrize`.
--
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]
