1"""
  2Timezone-related classes and functions.
  3"""
  4
  5from __future__ import annotations
  6
  7import functools
  8import zoneinfo
  9from contextlib import ContextDecorator
 10from contextvars import ContextVar
 11from datetime import UTC, datetime, time, timedelta, timezone, tzinfo
 12from types import TracebackType
 13
 14from plain.runtime import settings
 15
 16__all__ = [
 17    "get_fixed_timezone",
 18    "get_default_timezone",
 19    "get_default_timezone_name",
 20    "get_current_timezone",
 21    "get_current_timezone_name",
 22    "activate",
 23    "deactivate",
 24    "override",
 25    "localtime",
 26    "now",
 27    "is_aware",
 28    "is_naive",
 29    "make_aware",
 30    "make_naive",
 31]
 32
 33
 34def get_fixed_timezone(offset: int | timedelta) -> timezone:
 35    """Return a tzinfo instance with a fixed offset from UTC."""
 36    if isinstance(offset, timedelta):
 37        offset = int(offset.total_seconds() // 60)
 38    sign = "-" if offset < 0 else "+"
 39    hhmm = "%02d%02d" % divmod(abs(offset), 60)  # noqa: UP031
 40    name = sign + hhmm
 41    return timezone(timedelta(minutes=offset), name)
 42
 43
 44# In order to avoid accessing settings at compile time,
 45# wrap the logic in a function and cache the result.
 46@functools.lru_cache
 47def get_default_timezone() -> zoneinfo.ZoneInfo:
 48    """
 49    Return the default time zone as a tzinfo instance.
 50
 51    This is the time zone defined by settings.TIME_ZONE.
 52    """
 53    return zoneinfo.ZoneInfo(settings.TIME_ZONE)
 54
 55
 56# This function exists for consistency with get_current_timezone_name
 57def get_default_timezone_name() -> str:
 58    """Return the name of the default time zone."""
 59    return _get_timezone_name(get_default_timezone())
 60
 61
 62_active: ContextVar[tzinfo | None] = ContextVar("_active", default=None)
 63
 64
 65def get_current_timezone() -> tzinfo:
 66    """Return the currently active time zone as a tzinfo instance."""
 67    tz = _active.get()
 68    return tz if tz is not None else get_default_timezone()
 69
 70
 71def get_current_timezone_name() -> str:
 72    """Return the name of the currently active time zone."""
 73    return _get_timezone_name(get_current_timezone())
 74
 75
 76def _get_timezone_name(timezone: tzinfo) -> str:
 77    """
 78    Return the offset for fixed offset timezones, or the name of timezone if
 79    not set.
 80    """
 81    return timezone.tzname(None) or str(timezone)
 82
 83
 84# Timezone selection functions.
 85
 86# These functions don't change os.environ['TZ'] and call time.tzset()
 87# because it isn't thread safe.
 88
 89
 90def activate(timezone: tzinfo | str) -> None:
 91    """
 92    Set the time zone for the current context.
 93
 94    The ``timezone`` argument must be an instance of a tzinfo subclass or a
 95    time zone name.
 96    """
 97    if isinstance(timezone, tzinfo):
 98        _active.set(timezone)
 99    elif isinstance(timezone, str):
100        _active.set(zoneinfo.ZoneInfo(timezone))
101    else:
102        raise ValueError(f"Invalid timezone: {timezone!r}")
103
104
105def deactivate() -> None:
106    """
107    Unset the time zone for the current context.
108
109    Plain will then use the time zone defined by settings.TIME_ZONE.
110    """
111    _active.set(None)
112
113
114class override(ContextDecorator):
115    """
116    Temporarily set the time zone for the current context.
117
118    This is a context manager that uses plain.utils.timezone.activate()
119    to set the timezone on entry and restores the previously active timezone
120    on exit.
121
122    The ``timezone`` argument must be an instance of a ``tzinfo`` subclass, a
123    time zone name, or ``None``. If it is ``None``, Plain enables the default
124    time zone.
125    """
126
127    def __init__(self, timezone: tzinfo | str | None) -> None:
128        self.timezone = timezone
129
130    def __enter__(self) -> None:
131        if self.timezone is None:
132            self._token = _active.set(None)
133        elif isinstance(self.timezone, str):
134            self._token = _active.set(zoneinfo.ZoneInfo(self.timezone))
135        else:
136            self._token = _active.set(self.timezone)
137
138    def __exit__(
139        self,
140        exc_type: type[BaseException] | None,
141        exc_value: BaseException | None,
142        traceback: TracebackType | None,
143    ) -> None:
144        _active.reset(self._token)
145
146
147# Utilities
148
149
150def localtime(
151    value: datetime | None = None, timezone: tzinfo | None = None
152) -> datetime:
153    """
154    Convert an aware datetime.datetime to local time.
155
156    Only aware datetimes are allowed. When value is omitted, it defaults to
157    now().
158
159    Local time is defined by the current time zone, unless another time zone
160    is specified.
161    """
162    if value is None:
163        value = now()
164    if timezone is None:
165        timezone = get_current_timezone()
166    # Emulate the behavior of astimezone() on Python < 3.6.
167    if is_naive(value):
168        raise ValueError("localtime() cannot be applied to a naive datetime")
169    return value.astimezone(timezone)
170
171
172def now() -> datetime:
173    """
174    Return a timezone aware datetime.
175    """
176    return datetime.now(tz=UTC)
177
178
179# By design, these four functions don't perform any checks on their arguments.
180# The caller should ensure that they don't receive an invalid value like None.
181
182
183def is_aware(value: datetime | time) -> bool:
184    """
185    Determine if a given datetime.datetime or datetime.time is aware.
186
187    The concept is defined in Python's docs:
188    https://docs.python.org/library/datetime.html#datetime.tzinfo
189
190    Assuming value.tzinfo is either None or a proper datetime.tzinfo,
191    value.utcoffset() implements the appropriate logic.
192    """
193    return value.utcoffset() is not None
194
195
196def is_naive(value: datetime | time) -> bool:
197    """
198    Determine if a given datetime.datetime or datetime.time is naive.
199
200    The concept is defined in Python's docs:
201    https://docs.python.org/library/datetime.html#datetime.tzinfo
202
203    Assuming value.tzinfo is either None or a proper datetime.tzinfo,
204    value.utcoffset() implements the appropriate logic.
205    """
206    return value.utcoffset() is None
207
208
209def make_aware(value: datetime, timezone: tzinfo | None = None) -> datetime:
210    """Make a naive datetime.datetime in a given time zone aware."""
211    if timezone is None:
212        timezone = get_current_timezone()
213    # Check that we won't overwrite the timezone of an aware datetime.
214    if is_aware(value):
215        raise ValueError(f"make_aware expects a naive datetime, got {value}")
216    # This may be wrong around DST changes!
217    return value.replace(tzinfo=timezone)
218
219
220def make_naive(value: datetime, timezone: tzinfo | None = None) -> datetime:
221    """Make an aware datetime.datetime naive in a given time zone."""
222    if timezone is None:
223        timezone = get_current_timezone()
224    # Emulate the behavior of astimezone() on Python < 3.6.
225    if is_naive(value):
226        raise ValueError("make_naive() cannot be applied to a naive datetime")
227    return value.astimezone(timezone).replace(tzinfo=None)
228
229
230def _datetime_ambiguous_or_imaginary(dt: datetime, tz: tzinfo) -> bool:
231    return tz.utcoffset(dt.replace(fold=not dt.fold)) != tz.utcoffset(dt)