Skip to content

Commit

Permalink
Make pandas optional in workday calendar example (#30660)
Browse files Browse the repository at this point in the history
The workday calendar expected pandas to be available and it is part
of our examples, however Airflow does not have pandas as a core
dependency, so in case someone does not have pandas installed, importing
of the workday example would fail.

This change makes pandas optional and fallbacks to regular working
days for the example in case it is not available (including warning
about it). It also fixes a slight inefficiency where the
USFederalHoliday calendar has been created every time next workday
was calculated.
  • Loading branch information
potiuk committed Apr 15, 2023
1 parent d87eefc commit 5b42aa3
Showing 1 changed file with 16 additions and 6 deletions.
22 changes: 16 additions & 6 deletions airflow/example_dags/plugins/workday.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,28 +18,38 @@
"""Plugin to demonstrate timetable registration and accommodate example DAGs."""
from __future__ import annotations

import logging
from datetime import timedelta

# [START howto_timetable]
from pandas.tseries.holiday import USFederalHolidayCalendar
from pendulum import UTC, Date, DateTime, Time

from airflow.plugins_manager import AirflowPlugin
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable

log = logging.getLogger(__name__)
holiday_calendar = None

try:
from pandas.tseries.holiday import USFederalHolidayCalendar

holiday_calendar = USFederalHolidayCalendar()
except ImportError:
log.warning("Could not import pandas. Holidays will not be considered.")


class AfterWorkdayTimetable(Timetable):
def get_next_workday(self, d: DateTime, incr=1) -> DateTime:
cal = USFederalHolidayCalendar()
next_start = d
while True:
if next_start.weekday() in (5, 6): # If next start is in the weekend go to next day
next_start = next_start + incr * timedelta(days=1)
continue
holidays = cal.holidays(start=next_start, end=next_start).to_pydatetime()
if next_start in holidays: # If next start is a holiday go to next day
next_start = next_start + incr * timedelta(days=1)
continue
if holiday_calendar is not None:
holidays = holiday_calendar.holidays(start=next_start, end=next_start).to_pydatetime()
if next_start in holidays: # If next start is a holiday go to next day
next_start = next_start + incr * timedelta(days=1)
continue
break
return next_start

Expand Down

0 comments on commit 5b42aa3

Please sign in to comment.