From 1c689d7a52e945bb58964861af96cc9affb456bc Mon Sep 17 00:00:00 2001
From: Matthew Harrigan <mpharrigan@google.com>
Date: Thu, 21 Apr 2022 10:08:42 -0700
Subject: [PATCH] Serialize datetime (#5274)

spun off from #5152

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

cc @maffoo
---
 cirq/json_resolver_cache.py                   | 13 +++++++++-
 cirq/protocols/json_serialization.py          | 12 +++++++++
 cirq/protocols/json_serialization_test.py     | 26 +++++++++++++++++++
 .../json_test_data/datetime.datetime.json     |  4 +++
 .../json_test_data/datetime.datetime.repr     |  1 +
 5 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 cirq/protocols/json_test_data/datetime.datetime.json
 create mode 100644 cirq/protocols/json_test_data/datetime.datetime.repr

diff --git a/cirq/json_resolver_cache.py b/cirq/json_resolver_cache.py
index 80f1a259b49..656eb3fbebc 100644
--- a/cirq/json_resolver_cache.py
+++ b/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,
     }
diff --git a/cirq/protocols/json_serialization.py b/cirq/protocols/json_serialization.py
index e62f91eb062..ce514e4a956 100644
--- a/cirq/protocols/json_serialization.py
+++ b/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
 
 
diff --git a/cirq/protocols/json_serialization_test.py b/cirq/protocols/json_serialization_test.py
index b0edff1176f..e969e7f0f5f 100644
--- a/cirq/protocols/json_serialization_test.py
+++ b/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))
diff --git a/cirq/protocols/json_test_data/datetime.datetime.json b/cirq/protocols/json_test_data/datetime.datetime.json
new file mode 100644
index 00000000000..7c5f490b778
--- /dev/null
+++ b/cirq/protocols/json_test_data/datetime.datetime.json
@@ -0,0 +1,4 @@
+{
+  "cirq_type": "datetime.datetime",
+  "timestamp": 1648776225.0
+}
\ No newline at end of file
diff --git a/cirq/protocols/json_test_data/datetime.datetime.repr b/cirq/protocols/json_test_data/datetime.datetime.repr
new file mode 100644
index 00000000000..92f4bbde192
--- /dev/null
+++ b/cirq/protocols/json_test_data/datetime.datetime.repr
@@ -0,0 +1 @@
+datetime.datetime(2022, 4, 1, 1, 23, 45, tzinfo=datetime.timezone.utc)