1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
19
20import datetime as dt
21from importlib import metadata
22from typing import TYPE_CHECKING, overload
23
24import pendulum
25from dateutil.relativedelta import relativedelta
26from packaging import version
27from pendulum.datetime import DateTime
28
29if TYPE_CHECKING:
30 from pendulum.tz.timezone import FixedTimezone, Timezone
31
32 from airflow.typing_compat import Literal
33
34_PENDULUM3 = version.parse(metadata.version("pendulum")).major == 3
35# UTC Timezone as a tzinfo instance. Actual value depends on pendulum version:
36# - Timezone("UTC") in pendulum 3
37# - FixedTimezone(0, "UTC") in pendulum 2
38utc = pendulum.UTC
39
40
41def is_localized(value):
42 """Determine if a given datetime.datetime is aware.
43
44 The concept is defined in Python documentation. Assuming the tzinfo is
45 either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()``
46 implements the appropriate logic.
47
48 .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo
49 """
50 return value.utcoffset() is not None
51
52
53def is_naive(value):
54 """Determine if a given datetime.datetime is naive.
55
56 The concept is defined in Python documentation. Assuming the tzinfo is
57 either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()``
58 implements the appropriate logic.
59
60 .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo
61 """
62 return value.utcoffset() is None
63
64
65def utcnow() -> dt.datetime:
66 """Get the current date and time in UTC."""
67 return dt.datetime.now(tz=utc)
68
69
70def utc_epoch() -> dt.datetime:
71 """Get the epoch in the user's timezone."""
72 # pendulum utcnow() is not used as that sets a TimezoneInfo object
73 # instead of a Timezone. This is not picklable and also creates issues
74 # when using replace()
75 result = dt.datetime(1970, 1, 1)
76 result = result.replace(tzinfo=utc)
77
78 return result
79
80
81@overload
82def convert_to_utc(value: None) -> None: ...
83
84
85@overload
86def convert_to_utc(value: dt.datetime) -> DateTime: ...
87
88
89def convert_to_utc(value: dt.datetime | None) -> DateTime | None:
90 """Create a datetime with the default timezone added if none is associated.
91
92 :param value: datetime
93 :return: datetime with tzinfo
94 """
95 if value is None:
96 return value
97
98 if not is_localized(value):
99 from airflow.settings import TIMEZONE
100
101 value = pendulum.instance(value, TIMEZONE)
102
103 return pendulum.instance(value.astimezone(utc))
104
105
106@overload
107def make_aware(value: None, timezone: dt.tzinfo | None = None) -> None: ...
108
109
110@overload
111def make_aware(value: DateTime, timezone: dt.tzinfo | None = None) -> DateTime: ...
112
113
114@overload
115def make_aware(value: dt.datetime, timezone: dt.tzinfo | None = None) -> dt.datetime: ...
116
117
118def make_aware(value: dt.datetime | None, timezone: dt.tzinfo | None = None) -> dt.datetime | None:
119 """
120 Make a naive datetime.datetime in a given time zone aware.
121
122 :param value: datetime
123 :param timezone: timezone
124 :return: localized datetime in settings.TIMEZONE or timezone
125 """
126 if timezone is None:
127 from airflow.settings import TIMEZONE
128
129 timezone = TIMEZONE
130
131 if not value:
132 return None
133
134 # Check that we won't overwrite the timezone of an aware datetime.
135 if is_localized(value):
136 raise ValueError(f"make_aware expects a naive datetime, got {value}")
137 # In case we move clock back we want to schedule the run at the time of the second
138 # instance of the same clock time rather than the first one.
139 # Fold parameter has no impact in other cases, so we can safely set it to 1 here
140 value = value.replace(fold=1)
141 localized = getattr(timezone, "localize", None)
142 if localized is not None:
143 # This method is available for pytz time zones
144 return localized(value)
145 convert = getattr(timezone, "convert", None)
146 if convert is not None:
147 # For pendulum
148 return convert(value)
149 # This may be wrong around DST changes!
150 return value.replace(tzinfo=timezone)
151
152
153def make_naive(value, timezone=None):
154 """
155 Make an aware datetime.datetime naive in a given time zone.
156
157 :param value: datetime
158 :param timezone: timezone
159 :return: naive datetime
160 """
161 if timezone is None:
162 from airflow.settings import TIMEZONE
163
164 timezone = TIMEZONE
165
166 # Emulate the behavior of astimezone() on Python < 3.6.
167 if is_naive(value):
168 raise ValueError("make_naive() cannot be applied to a naive datetime")
169
170 date = value.astimezone(timezone)
171
172 # cross library compatibility
173 naive = dt.datetime(
174 date.year, date.month, date.day, date.hour, date.minute, date.second, date.microsecond
175 )
176
177 return naive
178
179
180def datetime(*args, **kwargs):
181 """
182 Wrap around datetime.datetime to add settings.TIMEZONE if tzinfo not specified.
183
184 :return: datetime.datetime
185 """
186 if "tzinfo" not in kwargs:
187 from airflow.settings import TIMEZONE
188
189 kwargs["tzinfo"] = TIMEZONE
190
191 return dt.datetime(*args, **kwargs)
192
193
194def parse(string: str, timezone=None, *, strict=False) -> DateTime:
195 """
196 Parse a time string and return an aware datetime.
197
198 :param string: time string
199 :param timezone: the timezone
200 :param strict: if False, it will fall back on the dateutil parser if unable to parse with pendulum
201 """
202 from airflow.settings import TIMEZONE
203
204 return pendulum.parse(string, tz=timezone or TIMEZONE, strict=strict) # type: ignore
205
206
207@overload
208def coerce_datetime(v: None, tz: dt.tzinfo | None = None) -> None: ...
209
210
211@overload
212def coerce_datetime(v: DateTime, tz: dt.tzinfo | None = None) -> DateTime: ...
213
214
215@overload
216def coerce_datetime(v: dt.datetime, tz: dt.tzinfo | None = None) -> DateTime: ...
217
218
219def coerce_datetime(v: dt.datetime | None, tz: dt.tzinfo | None = None) -> DateTime | None:
220 """Convert ``v`` into a timezone-aware ``pendulum.DateTime``.
221
222 * If ``v`` is *None*, *None* is returned.
223 * If ``v`` is a naive datetime, it is converted to an aware Pendulum DateTime.
224 * If ``v`` is an aware datetime, it is converted to a Pendulum DateTime.
225 Note that ``tz`` is **not** taken into account in this case; the datetime
226 will maintain its original tzinfo!
227 """
228 if v is None:
229 return None
230 if isinstance(v, DateTime):
231 return v if v.tzinfo else make_aware(v, tz)
232 # Only dt.datetime is left here.
233 return pendulum.instance(v if v.tzinfo else make_aware(v, tz))
234
235
236def td_format(td_object: None | dt.timedelta | float | int) -> str | None:
237 """
238 Format a timedelta object or float/int into a readable string for time duration.
239
240 For example timedelta(seconds=3752) would become `1h:2M:32s`.
241 If the time is less than a second, the return will be `<1s`.
242 """
243 if not td_object:
244 return None
245 if isinstance(td_object, dt.timedelta):
246 delta = relativedelta() + td_object
247 else:
248 delta = relativedelta(seconds=int(td_object))
249 # relativedelta for timedelta cannot convert days to months
250 # so calculate months by assuming 30 day months and normalize
251 months, delta.days = divmod(delta.days, 30)
252 delta = delta.normalized() + relativedelta(months=months)
253
254 def _format_part(key: str) -> str:
255 value = int(getattr(delta, key))
256 if value < 1:
257 return ""
258 # distinguish between month/minute following strftime format
259 # and take first char of each unit, i.e. years='y', days='d'
260 if key == "minutes":
261 key = key.upper()
262 key = key[0]
263 return f"{value}{key}"
264
265 parts = map(_format_part, ("years", "months", "days", "hours", "minutes", "seconds"))
266 joined = ":".join(part for part in parts if part)
267 if not joined:
268 return "<1s"
269 return joined
270
271
272def parse_timezone(name: str | int) -> FixedTimezone | Timezone:
273 """
274 Parse timezone and return one of the pendulum Timezone.
275
276 Provide the same interface as ``pendulum.timezone(name)``
277
278 :param name: Either IANA timezone or offset to UTC in seconds.
279
280 :meta private:
281 """
282 if _PENDULUM3:
283 # This only presented in pendulum 3 and code do not reached into the pendulum 2
284 return pendulum.timezone(name) # type: ignore[operator]
285 # In pendulum 2 this refers to the function, in pendulum 3 refers to the module
286 return pendulum.tz.timezone(name) # type: ignore[operator]
287
288
289def local_timezone() -> FixedTimezone | Timezone:
290 """
291 Return local timezone.
292
293 Provide the same interface as ``pendulum.tz.local_timezone()``
294
295 :meta private:
296 """
297 return pendulum.tz.local_timezone()
298
299
300def from_timestamp(
301 timestamp: int | float, tz: str | FixedTimezone | Timezone | Literal["local"] = utc
302) -> DateTime:
303 """
304 Parse timestamp and return DateTime in a given time zone.
305
306 :param timestamp: epoch time in seconds.
307 :param tz: In which timezone should return a resulting object.
308 Could be either one of pendulum timezone, IANA timezone or `local` literal.
309
310 :meta private:
311 """
312 result = coerce_datetime(dt.datetime.fromtimestamp(timestamp, tz=utc))
313 if tz != utc or tz != "UTC":
314 if isinstance(tz, str) and tz.lower() == "local":
315 tz = local_timezone()
316 result = result.in_timezone(tz)
317 return result