ext.py 5.5 KB

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