timezone.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321
  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.
  18. from __future__ import annotations
  19. import datetime as dt
  20. from importlib import metadata
  21. from typing import TYPE_CHECKING, overload
  22. import pendulum
  23. from dateutil.relativedelta import relativedelta
  24. from packaging import version
  25. from pendulum.datetime import DateTime
  26. if TYPE_CHECKING:
  27. from pendulum.tz.timezone import FixedTimezone, Timezone
  28. from airflow.typing_compat import Literal
  29. _PENDULUM3 = version.parse(metadata.version("pendulum")).major == 3
  30. # UTC Timezone as a tzinfo instance. Actual value depends on pendulum version:
  31. # - Timezone("UTC") in pendulum 3
  32. # - FixedTimezone(0, "UTC") in pendulum 2
  33. utc = pendulum.UTC
  34. def is_localized(value):
  35. """
  36. Determine if a given datetime.datetime is aware.
  37. The concept is defined in Python documentation. Assuming the tzinfo is
  38. either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()``
  39. implements the appropriate logic.
  40. .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo
  41. """
  42. return value.utcoffset() is not None
  43. def is_naive(value):
  44. """
  45. Determine if a given datetime.datetime is naive.
  46. The concept is defined in Python documentation. Assuming the tzinfo is
  47. either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()``
  48. implements the appropriate logic.
  49. .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo
  50. """
  51. return value.utcoffset() is None
  52. def utcnow() -> dt.datetime:
  53. """Get the current date and time in UTC."""
  54. return dt.datetime.now(tz=utc)
  55. def utc_epoch() -> dt.datetime:
  56. """Get the epoch in the user's timezone."""
  57. # pendulum utcnow() is not used as that sets a TimezoneInfo object
  58. # instead of a Timezone. This is not picklable and also creates issues
  59. # when using replace()
  60. result = dt.datetime(1970, 1, 1)
  61. result = result.replace(tzinfo=utc)
  62. return result
  63. @overload
  64. def convert_to_utc(value: None) -> None: ...
  65. @overload
  66. def convert_to_utc(value: dt.datetime) -> DateTime: ...
  67. def convert_to_utc(value: dt.datetime | None) -> DateTime | None:
  68. """
  69. Create a datetime with the default timezone added if none is associated.
  70. :param value: datetime
  71. :return: datetime with tzinfo
  72. """
  73. if value is None:
  74. return value
  75. if not is_localized(value):
  76. from airflow.settings import TIMEZONE
  77. value = pendulum.instance(value, TIMEZONE)
  78. return pendulum.instance(value.astimezone(utc))
  79. @overload
  80. def make_aware(value: None, timezone: dt.tzinfo | None = None) -> None: ...
  81. @overload
  82. def make_aware(value: DateTime, timezone: dt.tzinfo | None = None) -> DateTime: ...
  83. @overload
  84. def make_aware(value: dt.datetime, timezone: dt.tzinfo | None = None) -> dt.datetime: ...
  85. def make_aware(value: dt.datetime | None, timezone: dt.tzinfo | None = None) -> dt.datetime | None:
  86. """
  87. Make a naive datetime.datetime in a given time zone aware.
  88. :param value: datetime
  89. :param timezone: timezone
  90. :return: localized datetime in settings.TIMEZONE or timezone
  91. """
  92. if timezone is None:
  93. from airflow.settings import TIMEZONE
  94. timezone = TIMEZONE
  95. if not value:
  96. return None
  97. # Check that we won't overwrite the timezone of an aware datetime.
  98. if is_localized(value):
  99. raise ValueError(f"make_aware expects a naive datetime, got {value}")
  100. # In case we move clock back we want to schedule the run at the time of the second
  101. # instance of the same clock time rather than the first one.
  102. # Fold parameter has no impact in other cases, so we can safely set it to 1 here
  103. value = value.replace(fold=1)
  104. localized = getattr(timezone, "localize", None)
  105. if localized is not None:
  106. # This method is available for pytz time zones
  107. return localized(value)
  108. convert = getattr(timezone, "convert", None)
  109. if convert is not None:
  110. # For pendulum
  111. return convert(value)
  112. # This may be wrong around DST changes!
  113. return value.replace(tzinfo=timezone)
  114. def make_naive(value, timezone=None):
  115. """
  116. Make an aware datetime.datetime naive in a given time zone.
  117. :param value: datetime
  118. :param timezone: timezone
  119. :return: naive datetime
  120. """
  121. if timezone is None:
  122. from airflow.settings import TIMEZONE
  123. timezone = TIMEZONE
  124. # Emulate the behavior of astimezone() on Python < 3.6.
  125. if is_naive(value):
  126. raise ValueError("make_naive() cannot be applied to a naive datetime")
  127. date = value.astimezone(timezone)
  128. # cross library compatibility
  129. naive = dt.datetime(
  130. date.year, date.month, date.day, date.hour, date.minute, date.second, date.microsecond
  131. )
  132. return naive
  133. def datetime(*args, **kwargs):
  134. """
  135. Wrap around datetime.datetime to add settings.TIMEZONE if tzinfo not specified.
  136. :return: datetime.datetime
  137. """
  138. if "tzinfo" not in kwargs:
  139. from airflow.settings import TIMEZONE
  140. kwargs["tzinfo"] = TIMEZONE
  141. return dt.datetime(*args, **kwargs)
  142. def parse(string: str, timezone=None, *, strict=False) -> DateTime:
  143. """
  144. Parse a time string and return an aware datetime.
  145. :param string: time string
  146. :param timezone: the timezone
  147. :param strict: if False, it will fall back on the dateutil parser if unable to parse with pendulum
  148. """
  149. from airflow.settings import TIMEZONE
  150. return pendulum.parse(string, tz=timezone or TIMEZONE, strict=strict) # type: ignore
  151. @overload
  152. def coerce_datetime(v: None, tz: dt.tzinfo | None = None) -> None: ...
  153. @overload
  154. def coerce_datetime(v: DateTime, tz: dt.tzinfo | None = None) -> DateTime: ...
  155. @overload
  156. def coerce_datetime(v: dt.datetime, tz: dt.tzinfo | None = None) -> DateTime: ...
  157. def coerce_datetime(v: dt.datetime | None, tz: dt.tzinfo | None = None) -> DateTime | None:
  158. """
  159. Convert ``v`` into a timezone-aware ``pendulum.DateTime``.
  160. * If ``v`` is *None*, *None* is returned.
  161. * If ``v`` is a naive datetime, it is converted to an aware Pendulum DateTime.
  162. * If ``v`` is an aware datetime, it is converted to a Pendulum DateTime.
  163. Note that ``tz`` is **not** taken into account in this case; the datetime
  164. will maintain its original tzinfo!
  165. """
  166. if v is None:
  167. return None
  168. if isinstance(v, DateTime):
  169. return v if v.tzinfo else make_aware(v, tz)
  170. # Only dt.datetime is left here.
  171. return pendulum.instance(v if v.tzinfo else make_aware(v, tz))
  172. def td_format(td_object: None | dt.timedelta | float | int) -> str | None:
  173. """
  174. Format a timedelta object or float/int into a readable string for time duration.
  175. For example timedelta(seconds=3752) would become `1h:2M:32s`.
  176. If the time is less than a second, the return will be `<1s`.
  177. """
  178. if not td_object:
  179. return None
  180. if isinstance(td_object, dt.timedelta):
  181. delta = relativedelta() + td_object
  182. else:
  183. delta = relativedelta(seconds=int(td_object))
  184. # relativedelta for timedelta cannot convert days to months
  185. # so calculate months by assuming 30 day months and normalize
  186. months, delta.days = divmod(delta.days, 30)
  187. delta = delta.normalized() + relativedelta(months=months)
  188. def _format_part(key: str) -> str:
  189. value = int(getattr(delta, key))
  190. if value < 1:
  191. return ""
  192. # distinguish between month/minute following strftime format
  193. # and take first char of each unit, i.e. years='y', days='d'
  194. if key == "minutes":
  195. key = key.upper()
  196. key = key[0]
  197. return f"{value}{key}"
  198. parts = map(_format_part, ("years", "months", "days", "hours", "minutes", "seconds"))
  199. joined = ":".join(part for part in parts if part)
  200. if not joined:
  201. return "<1s"
  202. return joined
  203. def parse_timezone(name: str | int) -> FixedTimezone | Timezone:
  204. """
  205. Parse timezone and return one of the pendulum Timezone.
  206. Provide the same interface as ``pendulum.timezone(name)``
  207. :param name: Either IANA timezone or offset to UTC in seconds.
  208. :meta private:
  209. """
  210. if _PENDULUM3:
  211. # This only presented in pendulum 3 and code do not reached into the pendulum 2
  212. return pendulum.timezone(name) # type: ignore[operator]
  213. # In pendulum 2 this refers to the function, in pendulum 3 refers to the module
  214. return pendulum.tz.timezone(name) # type: ignore[operator]
  215. def local_timezone() -> FixedTimezone | Timezone:
  216. """
  217. Return local timezone.
  218. Provide the same interface as ``pendulum.tz.local_timezone()``
  219. :meta private:
  220. """
  221. return pendulum.tz.local_timezone()
  222. def from_timestamp(
  223. timestamp: int | float, tz: str | FixedTimezone | Timezone | Literal["local"] = utc
  224. ) -> DateTime:
  225. """
  226. Parse timestamp and return DateTime in a given time zone.
  227. :param timestamp: epoch time in seconds.
  228. :param tz: In which timezone should return a resulting object.
  229. Could be either one of pendulum timezone, IANA timezone or `local` literal.
  230. :meta private:
  231. """
  232. result = coerce_datetime(dt.datetime.fromtimestamp(timestamp, tz=utc))
  233. if tz != utc or tz != "UTC":
  234. if isinstance(tz, str) and tz.lower() == "local":
  235. tz = local_timezone()
  236. result = result.in_timezone(tz)
  237. return result