dabla commented on code in PR #75434:
URL: https://github.com/apache/airflow/pull/75434#discussion_r4514432735
##########
providers/standard/src/airflow/providers/standard/operators/weekday.py:
##########
@@ -111,10 +111,16 @@ def __init__(
def choose_branch(self, context: Context) -> str | Iterable[str]:
if self.use_task_logical_date:
- now = context.get("logical_date")
- if not now:
- dag_run = context.get("dag_run")
- now = dag_run.run_after # type: ignore[union-attr, assignment]
+ logical_date = context.get("logical_date")
+ dag_run = context.get("dag_run")
+
+ if not (logical_date or (dag_run and dag_run.run_after)):
+ raise ValueError(
+ "Either `logical_date` or `run_after` should be provided
in the task context when "
+ "`use_task_logical_date` is True"
+ )
+
+ now = logical_date if logical_date else dag_run.run_after # type:
ignore[union-attr]
Review Comment:
**[warning]** The same block is now in three places. Please move it into one
shared helper, which also removes the `type: ignore`.
With this PR, `DayOfWeekSensor.poke`,
`BranchDayOfWeekOperator.choose_branch` and
`BranchDateTimeOperator.choose_branch` all have the same nine lines and the
same error message, so any change to the message has to be made three times
([DRY](https://en.wikipedia.org/wiki/Don%27t_repeat_yourself)). Please add one
helper, for example in `providers/standard/utils/`, and call it from all three
so they stay in step. If it narrows with `:=` instead of a ternary, the `#
type: ignore[union-attr]` is no longer needed either:
```python
def logical_date_or_run_after(context: Context) -> datetime.datetime:
if logical_date := context.get("logical_date"):
return logical_date
if (dag_run := context.get("dag_run")) and dag_run.run_after:
return dag_run.run_after
raise ValueError(
"Either `logical_date` or `run_after` should be provided in the task
context when "
"`use_task_logical_date` is True"
)
```
The operators already have tests for both the fallback and the error, so the
refactor is covered.
--
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]
