ext.py 5.9 KB

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