dabla commented on code in PR #75434:
URL: https://github.com/apache/airflow/pull/75434#discussion_r4512952274
##########
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:
**[nit]** The same block is now in three places. A small shared helper could
replace them and drop 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)). One helper, for
example in `providers/standard/utils/`, used by all three would keep them 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"
)
```
Copying the sensor as it is also works, so this can go in a follow-up.
--
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]
