codekingpro/portable-devtools

1import datetime
2import struct
3from collections import namedtuple
4
5
6class ExtType(namedtuple("ExtType", "code data")):
7    """ExtType represents ext type in msgpack."""
8
9    def __new__(cls, code, data):
10        if not isinstance(code, int):
11            raise TypeError("code must be int")
12        if not isinstance(data, bytes):
13            raise TypeError("data must be bytes")
14        if not 0 <= code <= 127:
15            raise ValueError("code must be 0~127")
16        return super().__new__(cls, code, data)
17
18
19class Timestamp:
20    """Timestamp represents the Timestamp extension type in msgpack.
21
22    When built with Cython, msgpack uses C methods to pack and unpack `Timestamp`.
23    When using pure-Python msgpack, :func:`to_bytes` and :func:`from_bytes` are used to pack and
24    unpack `Timestamp`.
25
26    This class is immutable: Do not override seconds and nanoseconds.
27    """
28
29    __slots__ = ["seconds", "nanoseconds"]
30
31    def __init__(self, seconds, nanoseconds=0):
32        """Initialize a Timestamp object.
33
34        :param int seconds:
35            Number of seconds since the UNIX epoch (00:00:00 UTC Jan 1 1970, minus leap seconds).
36            May be negative.
37
38        :param int nanoseconds:
39            Number of nanoseconds to add to `seconds` to get fractional time.
40            Maximum is 999_999_999.  Default is 0.
41
42        Note: Negative times (before the UNIX epoch) are represented as neg. seconds + pos. ns.
43        """
44        if not isinstance(seconds, int):
45            raise TypeError("seconds must be an integer")
46        if not isinstance(nanoseconds, int):
47            raise TypeError("nanoseconds must be an integer")
48        if not (0 <= nanoseconds < 10**9):
49            raise ValueError("nanoseconds must be a non-negative integer less than 999999999.")
50        self.seconds = seconds
51        self.nanoseconds = nanoseconds
52
53    def __repr__(self):
54        """String representation of Timestamp."""
55        return f"Timestamp(seconds={self.seconds}, nanoseconds={self.nanoseconds})"
56
57    def __eq__(self, other):
58        """Check for equality with another Timestamp object"""
59        if type(other) is self.__class__:
60            return self.seconds == other.seconds and self.nanoseconds == other.nanoseconds
61        return False
62
63    def __ne__(self, other):
64        """not-equals method (see :func:`__eq__()`)"""
65        return not self.__eq__(other)
66
67    def __hash__(self):
68        return hash((self.seconds, self.nanoseconds))
69
70    @staticmethod
71    def from_bytes(b):
72        """Unpack bytes into a `Timestamp` object.
73
74        Used for pure-Python msgpack unpacking.
75
76        :param b: Payload from msgpack ext message with code -1
77        :type b: bytes
78
79        :returns: Timestamp object unpacked from msgpack ext payload
80        :rtype: Timestamp
81        """
82        if len(b) == 4:
83            seconds = struct.unpack("!L", b)[0]
84            nanoseconds = 0
85        elif len(b) == 8:
86            data64 = struct.unpack("!Q", b)[0]
87            seconds = data64 & 0x00000003FFFFFFFF
88            nanoseconds = data64 >> 34
89        elif len(b) == 12:
90            nanoseconds, seconds = struct.unpack("!Iq", b)
91        else:
92            raise ValueError(
93                "Timestamp type can only be created from 32, 64, or 96-bit byte objects"
94            )
95        return Timestamp(seconds, nanoseconds)
96
97    def to_bytes(self):
98        """Pack this Timestamp object into bytes.
99
100        Used for pure-Python msgpack packing.
101
102        :returns data: Payload for EXT message with code -1 (timestamp type)
103        :rtype: bytes
104        """
105        if (self.seconds >> 34) == 0:  # seconds is non-negative and fits in 34 bits
106            data64 = self.nanoseconds << 34 | self.seconds
107            if data64 & 0xFFFFFFFF00000000 == 0:
108                # nanoseconds is zero and seconds < 2**32, so timestamp 32
109                data = struct.pack("!L", data64)
110            else:
111                # timestamp 64
112                data = struct.pack("!Q", data64)
113        else:
114            # timestamp 96
115            data = struct.pack("!Iq", self.nanoseconds, self.seconds)
116        return data
117
118    @staticmethod
119    def from_unix(unix_sec):
120        """Create a Timestamp from posix timestamp in seconds.
121
122        :param unix_float: Posix timestamp in seconds.
123        :type unix_float: int or float
124        """
125        seconds = int(unix_sec // 1)
126        nanoseconds = int((unix_sec % 1) * 10**9)
127        return Timestamp(seconds, nanoseconds)
128
129    def to_unix(self):
130        """Get the timestamp as a floating-point value.
131
132        :returns: posix timestamp
133        :rtype: float
134        """
135        return self.seconds + self.nanoseconds / 1e9
136
137    @staticmethod
138    def from_unix_nano(unix_ns):
139        """Create a Timestamp from posix timestamp in nanoseconds.
140
141        :param int unix_ns: Posix timestamp in nanoseconds.
142        :rtype: Timestamp
143        """
144        return Timestamp(*divmod(unix_ns, 10**9))
145
146    def to_unix_nano(self):
147        """Get the timestamp as a unixtime in nanoseconds.
148
149        :returns: posix timestamp in nanoseconds
150        :rtype: int
151        """
152        return self.seconds * 10**9 + self.nanoseconds
153
154    def to_datetime(self):
155        """Get the timestamp as a UTC datetime.
156
157        :rtype: `datetime.datetime`
158        """
159        utc = datetime.timezone.utc
160        return datetime.datetime.fromtimestamp(0, utc) + datetime.timedelta(
161            seconds=self.seconds, microseconds=self.nanoseconds // 1000
162        )
163
164    @staticmethod
165    def from_datetime(dt):
166        """Create a Timestamp from datetime with tzinfo.
167
168        :rtype: Timestamp
169        """
170        return Timestamp(seconds=int(dt.timestamp()), nanoseconds=dt.microsecond * 1000)
171