You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
158 lines
5.9 KiB
158 lines
5.9 KiB
# Copyright 2022-present MongoDB, Inc. |
|
# |
|
# Licensed under the Apache License, Version 2.0 (the "License"); you |
|
# may not use this file except in compliance with the License. You |
|
# may obtain a copy of the License at |
|
# |
|
# http://www.apache.org/licenses/LICENSE-2.0 |
|
# |
|
# Unless required by applicable law or agreed to in writing, software |
|
# distributed under the License is distributed on an "AS IS" BASIS, |
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
|
# implied. See the License for the specific language governing |
|
# permissions and limitations under the License. |
|
|
|
"""Tools for representing the BSON datetime type. |
|
|
|
.. versionadded:: 4.3 |
|
""" |
|
|
|
import calendar |
|
import datetime |
|
import functools |
|
from typing import Any, Union, cast |
|
|
|
from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions, DatetimeConversion |
|
from bson.tz_util import utc |
|
|
|
EPOCH_AWARE = datetime.datetime.fromtimestamp(0, utc) |
|
EPOCH_NAIVE = EPOCH_AWARE.replace(tzinfo=None) |
|
|
|
|
|
class DatetimeMS: |
|
"""Represents a BSON UTC datetime.""" |
|
|
|
__slots__ = ("_value",) |
|
|
|
def __init__(self, value: Union[int, datetime.datetime]): |
|
"""Represents a BSON UTC datetime. |
|
|
|
BSON UTC datetimes are defined as an int64 of milliseconds since the |
|
Unix epoch. The principal use of DatetimeMS is to represent |
|
datetimes outside the range of the Python builtin |
|
:class:`~datetime.datetime` class when |
|
encoding/decoding BSON. |
|
|
|
To decode UTC datetimes as a ``DatetimeMS``, `datetime_conversion` in |
|
:class:`~bson.CodecOptions` must be set to 'datetime_ms' or |
|
'datetime_auto'. See :ref:`handling-out-of-range-datetimes` for |
|
details. |
|
|
|
:Parameters: |
|
- `value`: An instance of :class:`datetime.datetime` to be |
|
represented as milliseconds since the Unix epoch, or int of |
|
milliseconds since the Unix epoch. |
|
""" |
|
if isinstance(value, int): |
|
if not (-(2**63) <= value <= 2**63 - 1): |
|
raise OverflowError("Must be a 64-bit integer of milliseconds") |
|
self._value = value |
|
elif isinstance(value, datetime.datetime): |
|
self._value = _datetime_to_millis(value) |
|
else: |
|
raise TypeError(f"{type(value)} is not a valid type for DatetimeMS") |
|
|
|
def __hash__(self) -> int: |
|
return hash(self._value) |
|
|
|
def __repr__(self) -> str: |
|
return type(self).__name__ + "(" + str(self._value) + ")" |
|
|
|
def __lt__(self, other: Union["DatetimeMS", int]) -> bool: |
|
return self._value < other |
|
|
|
def __le__(self, other: Union["DatetimeMS", int]) -> bool: |
|
return self._value <= other |
|
|
|
def __eq__(self, other: Any) -> bool: |
|
if isinstance(other, DatetimeMS): |
|
return self._value == other._value |
|
return False |
|
|
|
def __ne__(self, other: Any) -> bool: |
|
if isinstance(other, DatetimeMS): |
|
return self._value != other._value |
|
return True |
|
|
|
def __gt__(self, other: Union["DatetimeMS", int]) -> bool: |
|
return self._value > other |
|
|
|
def __ge__(self, other: Union["DatetimeMS", int]) -> bool: |
|
return self._value >= other |
|
|
|
_type_marker = 9 |
|
|
|
def as_datetime(self, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> datetime.datetime: |
|
"""Create a Python :class:`~datetime.datetime` from this DatetimeMS object. |
|
|
|
:Parameters: |
|
- `codec_options`: A CodecOptions instance for specifying how the |
|
resulting DatetimeMS object will be formatted using ``tz_aware`` |
|
and ``tz_info``. Defaults to |
|
:const:`~bson.codec_options.DEFAULT_CODEC_OPTIONS`. |
|
""" |
|
return cast(datetime.datetime, _millis_to_datetime(self._value, codec_options)) |
|
|
|
def __int__(self) -> int: |
|
return self._value |
|
|
|
|
|
# Inclusive and exclusive min and max for timezones. |
|
# Timezones are hashed by their offset, which is a timedelta |
|
# and therefore there are more than 24 possible timezones. |
|
@functools.lru_cache(maxsize=None) |
|
def _min_datetime_ms(tz: datetime.timezone = datetime.timezone.utc) -> int: |
|
return _datetime_to_millis(datetime.datetime.min.replace(tzinfo=tz)) |
|
|
|
|
|
@functools.lru_cache(maxsize=None) |
|
def _max_datetime_ms(tz: datetime.timezone = datetime.timezone.utc) -> int: |
|
return _datetime_to_millis(datetime.datetime.max.replace(tzinfo=tz)) |
|
|
|
|
|
def _millis_to_datetime(millis: int, opts: CodecOptions) -> Union[datetime.datetime, DatetimeMS]: |
|
"""Convert milliseconds since epoch UTC to datetime.""" |
|
if ( |
|
opts.datetime_conversion == DatetimeConversion.DATETIME |
|
or opts.datetime_conversion == DatetimeConversion.DATETIME_CLAMP |
|
or opts.datetime_conversion == DatetimeConversion.DATETIME_AUTO |
|
): |
|
tz = opts.tzinfo or datetime.timezone.utc |
|
if opts.datetime_conversion == DatetimeConversion.DATETIME_CLAMP: |
|
millis = max(_min_datetime_ms(tz), min(millis, _max_datetime_ms(tz))) |
|
elif opts.datetime_conversion == DatetimeConversion.DATETIME_AUTO: |
|
if not (_min_datetime_ms(tz) <= millis <= _max_datetime_ms(tz)): |
|
return DatetimeMS(millis) |
|
|
|
diff = ((millis % 1000) + 1000) % 1000 |
|
seconds = (millis - diff) // 1000 |
|
micros = diff * 1000 |
|
|
|
if opts.tz_aware: |
|
dt = EPOCH_AWARE + datetime.timedelta(seconds=seconds, microseconds=micros) |
|
if opts.tzinfo: |
|
dt = dt.astimezone(tz) |
|
return dt |
|
else: |
|
return EPOCH_NAIVE + datetime.timedelta(seconds=seconds, microseconds=micros) |
|
elif opts.datetime_conversion == DatetimeConversion.DATETIME_MS: |
|
return DatetimeMS(millis) |
|
else: |
|
raise ValueError("datetime_conversion must be an element of DatetimeConversion") |
|
|
|
|
|
def _datetime_to_millis(dtm: datetime.datetime) -> int: |
|
"""Convert datetime to milliseconds since epoch UTC.""" |
|
if dtm.utcoffset() is not None: |
|
dtm = dtm - dtm.utcoffset() # type: ignore |
|
return int(calendar.timegm(dtm.timetuple()) * 1000 + dtm.microsecond // 1000)
|
|
|