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