Serialize datetime (quantumlib#5274)

spun off from quantumlib#5152 

It's super important when running experiments to keep track of when things happen.

cc @maffoo

cirq/json_resolver_cache.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Methods for resolving JSON types during serialization."""
-
+import datetime
 import functools
 from typing import Dict, TYPE_CHECKING
 
@@ -54,6 +54,16 @@ def two_qubit_matrix_gate(matrix):
     def _parallel_gate_op(gate, qubits):
         return cirq.parallel_gate_op(gate, *qubits)
 
+    def _datetime(timestamp: float) -> datetime.datetime:
+        # As part of our serialization logic, we make sure we only serialize "aware"
+        # datetimes with the UTC timezone, so we implicitly add back in the UTC timezone here.
+        #
+        # Please note: even if the assumption is somehow violated, the fact that we use
+        # unix timestamps should mean that the deserialized datetime should refer to the
+        # same point in time but may not satisfy o = read_json(to_json(o)) because the actual
+        # timezones, and hour fields will not be identical.
+        return datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc)
+
     import sympy
 
     return {
@@ -216,4 +226,5 @@ def _parallel_gate_op(gate, qubits):
         'sympy.E': lambda: sympy.E,
         'sympy.EulerGamma': lambda: sympy.EulerGamma,
         'complex': complex,
+        'datetime.datetime': _datetime,
     }

cirq/protocols/json_serialization.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import dataclasses
+import datetime
 import gzip
 import json
 import numbers
@@ -386,6 +387,17 @@ def default(self, o):
                 'index': o.index,
             }
 
+        # datetime
+        if isinstance(o, datetime.datetime):
+            if o.tzinfo is None or o.tzinfo.utcoffset(o) is None:
+                # Otherwise, the deserialized object may change depending on local timezone.
+                raise TypeError(
+                    "Can only serialize 'aware' datetime objects with `tzinfo`. "
+                    "Consider using e.g. `datetime.datetime.now(tz=datetime.timezone.utc)`"
+                )
+
+            return {'cirq_type': 'datetime.datetime', 'timestamp': o.timestamp()}
+
         return super().default(o)  # coverage: ignore
 
 

cirq/protocols/json_serialization_test.py

@@ -933,3 +933,29 @@ def test_numpy_values():
   "value": 1
 }"""
     )
+
+
+def test_basic_time_assertions():
+    naive_dt = datetime.datetime.now()
+    utc_dt = naive_dt.astimezone(datetime.timezone.utc)
+    assert naive_dt.timestamp() == utc_dt.timestamp()
+
+    re_utc = datetime.datetime.fromtimestamp(utc_dt.timestamp())
+    re_naive = datetime.datetime.fromtimestamp(naive_dt.timestamp())
+
+    assert re_utc == re_naive, 'roundtripping w/o tz turns to naive utc'
+    assert re_utc != utc_dt, 'roundtripping loses tzinfo'
+    assert naive_dt == re_naive, 'works, as long as you called fromtimestamp from the same timezone'
+
+
+def test_datetime():
+    naive_dt = datetime.datetime.now()
+
+    with pytest.raises(TypeError):
+        cirq.to_json(naive_dt)
+
+    utc_dt = naive_dt.astimezone(datetime.timezone.utc)
+    assert utc_dt == cirq.read_json(json_text=cirq.to_json(utc_dt))
+
+    pst_dt = naive_dt.astimezone(tz=datetime.timezone(offset=datetime.timedelta(hours=-8)))
+    assert utc_dt == cirq.read_json(json_text=cirq.to_json(pst_dt))

cirq/protocols/json_test_data/datetime.datetime.json

@@ -0,0 +1,4 @@
+{
+  "cirq_type": "datetime.datetime",
+  "timestamp": 1648776225.0
+}

cirq/protocols/json_test_data/datetime.datetime.repr

@@ -0,0 +1 @@
+datetime.datetime(2022, 4, 1, 1, 23, 45, tzinfo=datetime.timezone.utc)