1# Copyright 2017 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Helpers for :mod:`datetime`."""
16
17import datetime
18import re
19
20from google.protobuf import timestamp_pb2
21
22_UTC_EPOCH = datetime.datetime.fromtimestamp(0, datetime.timezone.utc)
23
24_RFC3339_MICROS = "%Y-%m-%dT%H:%M:%S.%fZ"
25_RFC3339_NO_FRACTION = "%Y-%m-%dT%H:%M:%S"
26# datetime.strptime cannot handle nanosecond precision: parse w/ regex
27_RFC3339_NANOS = re.compile(
28 r"""
29 (?P<no_fraction>
30 \d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2} # YYYY-MM-DDTHH:MM:SS
31 )
32 ( # Optional decimal part
33 \. # decimal point
34 (?P<nanos>\d{1,9}) # nanoseconds, maybe truncated
35 )?
36 Z # Zulu
37""",
38 re.VERBOSE,
39)
40
41
42def _from_microseconds(value):
43 """Convert timestamp in microseconds since the unix epoch to datetime.
44
45 Args:
46 value (float): The timestamp to convert, in microseconds.
47
48 Returns:
49 datetime.datetime: The datetime object equivalent to the timestamp in
50 UTC.
51 """
52 return _UTC_EPOCH + datetime.timedelta(microseconds=value)
53
54
55def _to_rfc3339(value, ignore_zone=True):
56 """Convert a datetime to an RFC3339 timestamp string.
57
58 Args:
59 value (datetime.datetime):
60 The datetime object to be converted to a string.
61 ignore_zone (bool): If True, then the timezone (if any) of the
62 datetime object is ignored and the datetime is treated as UTC.
63
64 Returns:
65 str: The RFC3339 formatted string representing the datetime.
66 """
67 if not ignore_zone and value.tzinfo is not None:
68 # Convert to UTC and remove the time zone info.
69 value = value.replace(tzinfo=None) - value.utcoffset()
70
71 return value.strftime(_RFC3339_MICROS)
72
73
74class DatetimeWithNanoseconds(datetime.datetime):
75 """Track nanosecond in addition to normal datetime attrs.
76
77 Nanosecond can be passed only as a keyword argument.
78 """
79
80 __slots__ = ("_nanosecond",)
81
82 # pylint: disable=arguments-differ
83 def __new__(cls, *args, **kw):
84 nanos = kw.pop("nanosecond", 0)
85 if nanos > 0:
86 if "microsecond" in kw:
87 raise TypeError("Specify only one of 'microsecond' or 'nanosecond'")
88 kw["microsecond"] = nanos // 1000
89 inst = datetime.datetime.__new__(cls, *args, **kw)
90 inst._nanosecond = nanos or 0
91 return inst
92
93 # pylint: disable=arguments-differ
94 def replace(self, *args, **kw):
95 """Return a date with the same value, except for those parameters given
96 new values by whichever keyword arguments are specified. For example,
97 if d == date(2002, 12, 31), then
98 d.replace(day=26) == date(2002, 12, 26).
99 NOTE: nanosecond and microsecond are mutually exclusive arguments.
100 """
101
102 ms_provided = "microsecond" in kw
103 ns_provided = "nanosecond" in kw
104 provided_ns = kw.pop("nanosecond", 0)
105
106 prev_nanos = self.nanosecond
107
108 if ms_provided and ns_provided:
109 raise TypeError("Specify only one of 'microsecond' or 'nanosecond'")
110
111 if ns_provided:
112 # if nanos were provided, manipulate microsecond kw arg to super
113 kw["microsecond"] = provided_ns // 1000
114 inst = super().replace(*args, **kw)
115
116 if ms_provided:
117 # ms were provided, nanos are invalid, build from ms
118 inst._nanosecond = inst.microsecond * 1000
119 elif ns_provided:
120 # ns were provided, replace nanoseconds to match after calling super
121 inst._nanosecond = provided_ns
122 else:
123 # if neither ms or ns were provided, passthru previous nanos.
124 inst._nanosecond = prev_nanos
125
126 return inst
127
128 @property
129 def nanosecond(self):
130 """Read-only: nanosecond precision."""
131 return self._nanosecond or self.microsecond * 1000
132
133 def rfc3339(self):
134 """Return an RFC3339-compliant timestamp.
135
136 Returns:
137 (str): Timestamp string according to RFC3339 spec.
138 """
139 if self._nanosecond == 0:
140 return _to_rfc3339(self)
141 nanos = str(self._nanosecond).rjust(9, "0").rstrip("0")
142 return "{}.{}Z".format(self.strftime(_RFC3339_NO_FRACTION), nanos)
143
144 @classmethod
145 def from_rfc3339(cls, stamp):
146 """Parse RFC3339-compliant timestamp, preserving nanoseconds.
147
148 Args:
149 stamp (str): RFC3339 stamp, with up to nanosecond precision
150
151 Returns:
152 :class:`DatetimeWithNanoseconds`:
153 an instance matching the timestamp string
154
155 Raises:
156 ValueError: if `stamp` does not match the expected format
157 """
158 with_nanos = _RFC3339_NANOS.match(stamp)
159 if with_nanos is None:
160 raise ValueError(
161 "Timestamp: {}, does not match pattern: {}".format(
162 stamp, _RFC3339_NANOS.pattern
163 )
164 )
165 bare = datetime.datetime.strptime(
166 with_nanos.group("no_fraction"), _RFC3339_NO_FRACTION
167 )
168 fraction = with_nanos.group("nanos")
169 if fraction is None:
170 nanos = 0
171 else:
172 scale = 9 - len(fraction)
173 nanos = int(fraction) * (10**scale)
174 return cls(
175 bare.year,
176 bare.month,
177 bare.day,
178 bare.hour,
179 bare.minute,
180 bare.second,
181 nanosecond=nanos,
182 tzinfo=datetime.timezone.utc,
183 )
184
185 def timestamp_pb(self):
186 """Return a timestamp message.
187
188 Returns:
189 (:class:`~google.protobuf.timestamp_pb2.Timestamp`): Timestamp message
190 """
191 inst = (
192 self
193 if self.tzinfo is not None
194 else self.replace(tzinfo=datetime.timezone.utc)
195 )
196 delta = inst - _UTC_EPOCH
197 seconds = int(delta.total_seconds())
198 nanos = self._nanosecond or self.microsecond * 1000
199 return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
200
201 @classmethod
202 def from_timestamp_pb(cls, stamp):
203 """Parse RFC3339-compliant timestamp, preserving nanoseconds.
204
205 Args:
206 stamp (:class:`~google.protobuf.timestamp_pb2.Timestamp`): timestamp message
207
208 Returns:
209 :class:`DatetimeWithNanoseconds`:
210 an instance matching the timestamp message
211 """
212 microseconds = int(stamp.seconds * 1e6)
213 bare = _from_microseconds(microseconds)
214 return cls(
215 bare.year,
216 bare.month,
217 bare.day,
218 bare.hour,
219 bare.minute,
220 bare.second,
221 nanosecond=stamp.nanos,
222 tzinfo=datetime.timezone.utc,
223 )