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