Add concise serialization tools.

cirq/protocols/json_serialization.py

@@ -408,10 +408,20 @@ def default(self, o):
         return super().default(o)  # coverage: ignore
 
 
-def _cirq_object_hook(d, resolvers: Sequence[JsonResolver]):
+def _cirq_object_hook(d, resolvers: Sequence[JsonResolver], context_map: Dict[str, Any]):
     if 'cirq_type' not in d:
         return d
 
+    if d['cirq_type'] == '_SerializedKey':
+        return _SerializedKey.read_from_context(context_map, **d)
+
+    if d['cirq_type'] == '_SerializedContext':
+        _SerializedContext.update_context(context_map, **d)
+        return None
+
+    if d['cirq_type'] == '_ContextualSerialization':
+        return _ContextualSerialization.deserialize_with_context(**d)
+
     for resolver in resolvers:
         cls = resolver(d['cirq_type'])
         if cls is not None:
@@ -429,6 +439,142 @@ def _cirq_object_hook(d, resolvers: Sequence[JsonResolver]):
     return cls(**d)
 
 
+class SerializableByKey(SupportsJSON):
+    """Protocol for objects that can be serialized to a key + context."""
+
+    @doc_private
+    def _serialization_key_(self) -> str:
+        """Returns a unique string identifier for this object.
+
+        This should only return the same value for two objects if they are
+        equal; otherwise, an error will occur if both are serialized into the
+        same JSON string.
+        """
+
+
+class _SerializedKey(SupportsJSON):
+    """Internal object for holding a SerializableByKey key.
+
+    This is a private type used in contextual serialization. Its deserialization
+    is context-dependent, and is not expected to match the original; in other
+    words, `cls._from_json_dict_(obj._json_dict_())` does not return
+    the original `obj` for this type.
+    """
+
+    def __init__(self, obj: SerializableByKey):
+        self.key = obj._serialization_key_()
+
+    def _json_dict_(self):
+        return obj_to_dict_helper(self, ['key'])
+
+    @classmethod
+    def _from_json_dict_(cls, **kwargs):
+        raise TypeError(f'Internal error: {cls} should never deserialize with _from_json_dict_.')
+
+    @classmethod
+    def read_from_context(cls, context_map, key, **kwargs):
+        return context_map[key]
+
+
+class _SerializedContext(SupportsJSON):
+    """Internal object for a single SerializableByKey key-to-object mapping.
+
+    This is a private type used in contextual serialization. Its deserialization
+    is context-dependent, and is not expected to match the original; in other
+    words, `cls._from_json_dict_(obj._json_dict_())` does not return
+    the original `obj` for this type.
+    """
+
+    def __init__(self, obj: SerializableByKey):
+        self.key = obj._serialization_key_()
+        self.obj = obj
+
+    def _json_dict_(self):
+        return obj_to_dict_helper(self, ['key', 'obj'])
+
+    @classmethod
+    def _from_json_dict_(cls, **kwargs):
+        raise TypeError(f'Internal error: {cls} should never deserialize with _from_json_dict_.')
+
+    @classmethod
+    def update_context(cls, context_map, key, obj, **kwargs):
+        context_map.update({key: obj})
+
+
+class _ContextualSerialization(SupportsJSON):
+    """Internal object for serializing an object with its context.
+
+    This is a private type used in contextual serialization. Its deserialization
+    is context-dependent, and is not expected to match the original; in other
+    words, `cls._from_json_dict_(obj._json_dict_())` does not return
+    the original `obj` for this type.
+    """
+
+    def __init__(self, obj: Any):
+        # Context information and the wrapped object are stored together in
+        # `object_dag` to ensure consistent serialization ordering.
+        self.object_dag = []
+        context_keys = set()
+        for sbk in get_serializable_by_keys(obj):
+            new_sc = _SerializedContext(sbk)
+            if new_sc.key not in context_keys:
+                self.object_dag.append(new_sc)
+                context_keys.add(new_sc.key)
+        self.object_dag += [obj]
+
+    def _json_dict_(self):
+        return obj_to_dict_helper(self, ['object_dag'])
+
+    @classmethod
+    def _from_json_dict_(cls, **kwargs):
+        raise TypeError(f'Internal error: {cls} should never deserialize with _from_json_dict_.')
+
+    @classmethod
+    def deserialize_with_context(cls, object_dag, **kwargs):
+        # The last element of object_dag is the object to be deserialized.
+        return object_dag[-1]
+
+
+def has_serializable_by_keys(obj: Any) -> bool:
+    """Returns true if obj contains one or more SerializableByKey objects."""
+    if hasattr(obj, '_serialization_key_'):
+        return True
+    json_dict = getattr(obj, '_json_dict_', lambda: None)()
+    if isinstance(json_dict, Dict):
+        return any(has_serializable_by_keys(v) for v in json_dict.values())
+
+    # Handle primitive container types.
+    if isinstance(obj, Dict):
+        return any(has_serializable_by_keys(elem) for pair in obj.items() for elem in pair)
+    if hasattr(obj, '__iter__') and not isinstance(obj, str):
+        return any(has_serializable_by_keys(elem) for elem in obj)
+    return False
+
+
+def get_serializable_by_keys(obj: Any) -> List[SerializableByKey]:
+    """Returns all SerializableByKeys contained by obj.
+
+    Objects are ordered such that nested objects appear before the object they
+    are nested inside. This is required to ensure
+    """
+    result = []
+    if hasattr(obj, '_serialization_key_'):
+        result.append(obj)
+    json_dict = getattr(obj, '_json_dict_', lambda: None)()
+    if isinstance(json_dict, Dict):
+        for v in json_dict.values():
+            result = get_serializable_by_keys(v) + result
+    if result:
+        return result
+
+    # Handle primitive container types.
+    if isinstance(obj, Dict):
+        return [sbk for pair in obj.items() for sbk in get_serializable_by_keys(pair)]
+    if hasattr(obj, '__iter__') and not isinstance(obj, str):
+        return [sbk for v in obj for sbk in get_serializable_by_keys(v)]
+    return []
+
+
 # pylint: disable=function-redefined
 @overload
 def to_json(
@@ -468,6 +614,32 @@ def to_json(
             party classes, prefer adding the _json_dict_ magic method
             to your classes rather than overriding this default.
     """
+    if has_serializable_by_keys(obj):
+
+        class ContextualEncoder(cls):  # type: ignore
+            """An encoder with a context map for concise serialization."""
+
+            # This map is populated gradually during serialization. An object
+            # with components defined in this map will represent those
+            # components using their keys instead of inline definition.
+            context_map: Dict[str, 'SerializableByKey'] = {}
+
+            def default(self, o):
+                skey = getattr(o, '_serialization_key_', lambda: None)()
+                if skey in ContextualEncoder.context_map:
+                    if ContextualEncoder.context_map[skey] == o._json_dict_():
+                        return _SerializedKey(o)._json_dict_()
+                    raise ValueError(
+                        'Found different objects with the same serialization key:'
+                        f'\n{ContextualEncoder.context_map[skey]}\n{o}'
+                    )
+                if skey is not None:
+                    ContextualEncoder.context_map[skey] = o._json_dict_()
+                return super().default(o)
+
+        obj = _ContextualSerialization(obj)
+        cls = ContextualEncoder
+
     if file_or_fn is None:
         return json.dumps(obj, indent=indent, cls=cls)
 
@@ -513,8 +685,10 @@ def read_json(
     if resolvers is None:
         resolvers = DEFAULT_RESOLVERS
 
+    context_map: Dict[str, 'SerializableByKey'] = {}
+
     def obj_hook(x):
-        return _cirq_object_hook(x, resolvers)
+        return _cirq_object_hook(x, resolvers, context_map)
 
     if json_text is not None:
         return json.loads(json_text, object_hook=obj_hook)

cirq/protocols/json_serialization_test.py

@@ -19,7 +19,7 @@
 import os
 import pathlib
 import textwrap
-from typing import Any, Iterator, List, Optional, Set, Tuple, Type
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Type
 
 import pytest
 
@@ -433,6 +433,115 @@ def test_sympy():
     assert_json_roundtrip_works(4 * t + 3 * s + 2)
 
 
+class SBKImpl:
+    """A test implementation of SerializableByKey."""
+
+    def __init__(
+        self,
+        name: str,
+        data_list: Optional[List] = None,
+        data_tuple: Optional[Tuple] = None,
+        data_dict: Optional[Dict] = None,
+    ):
+        self.name = name
+        self.data_list = data_list or []
+        self.data_tuple = data_tuple or ()
+        self.data_dict = data_dict or {}
+
+    def __eq__(self, other):
+        if not isinstance(other, SBKImpl):
+            return False
+        return (
+            self.name == other.name
+            and self.data_list == other.data_list
+            and self.data_tuple == other.data_tuple
+            and self.data_dict == other.data_dict
+        )
+
+    def _json_dict_(self):
+        return {
+            "cirq_type": "SBKImpl",
+            "name": self.name,
+            "data_list": self.data_list,
+            "data_tuple": self.data_tuple,
+            "data_dict": self.data_dict,
+        }
+
+    def _serialization_key_(self):
+        return self.name
+
+    @classmethod
+    def _from_json_dict_(cls, name, data_list, data_tuple, data_dict, **kwargs):
+        return cls(name, data_list, tuple(data_tuple), data_dict)
+
+
+def test_context_serialization():
+    def custom_resolver(name):
+        if name == 'SBKImpl':
+            return SBKImpl
+
+    test_resolvers = [custom_resolver] + cirq.DEFAULT_RESOLVERS
+
+    sbki_empty = SBKImpl('sbki_empty')
+    assert_json_roundtrip_works(sbki_empty, resolvers=test_resolvers)
+
+    sbki_list = SBKImpl('sbki_list', data_list=[sbki_empty, sbki_empty])
+    assert_json_roundtrip_works(sbki_list, resolvers=test_resolvers)
+
+    sbki_tuple = SBKImpl('sbki_tuple', data_tuple=(sbki_list, sbki_list))
+    assert_json_roundtrip_works(sbki_tuple, resolvers=test_resolvers)
+
+    sbki_dict = SBKImpl('sbki_dict', data_dict={'a': sbki_tuple, 'b': sbki_tuple})
+    assert_json_roundtrip_works(sbki_dict, resolvers=test_resolvers)
+
+    sbki_json = str(cirq.to_json(sbki_dict))
+    # There should be exactly one context item for each previous SBKImpl.
+    assert sbki_json.count('"cirq_type": "_SerializedContext"') == 4
+    # There should be exactly two key items for each of sbki_(empty|list|tuple),
+    # plus one for the top-level sbki_dict.
+    assert sbki_json.count('"cirq_type": "_SerializedKey"') == 7
+    # The final object should be a _SerializedKey for sbki_dict.
+    final_obj_idx = sbki_json.rfind('{')
+    final_obj = sbki_json[final_obj_idx : sbki_json.find('}', final_obj_idx) + 1]
+    assert (
+        final_obj
+        == """{
+      "cirq_type": "_SerializedKey",
+      "key": "sbki_dict"
+    }"""
+    )
+
+    list_sbki = [sbki_dict]
+    assert_json_roundtrip_works(list_sbki, resolvers=test_resolvers)
+
+    dict_sbki = {'a': sbki_dict}
+    assert_json_roundtrip_works(dict_sbki, resolvers=test_resolvers)
+
+    assert sbki_list != json_serialization._SerializedKey(sbki_list)
+    sbki_other_list = SBKImpl('sbki_list', data_list=[sbki_list])
+    with pytest.raises(ValueError, match='different objects with the same serialization key'):
+        _ = cirq.to_json(sbki_other_list)
+
+
+def test_internal_serializer_types():
+    sbki = SBKImpl('test_key')
+    test_key = json_serialization._SerializedKey(sbki)
+    test_context = json_serialization._SerializedContext(sbki)
+    test_serialization = json_serialization._ContextualSerialization(sbki)
+
+    key_json = test_key._json_dict_()
+    with pytest.raises(TypeError, match='_from_json_dict_'):
+        _ = json_serialization._SerializedKey._from_json_dict_(**key_json)
+
+    context_json = test_context._json_dict_()
+    with pytest.raises(TypeError, match='_from_json_dict_'):
+        _ = json_serialization._SerializedContext._from_json_dict_(**context_json)
+
+    serialization_json = test_serialization._json_dict_()
+    with pytest.raises(TypeError, match='_from_json_dict_'):
+        _ = json_serialization._ContextualSerialization._from_json_dict_(**serialization_json)
+
+
 def _write_test_data(key: str, *test_instances: Any):
     """Helper method for creating initial test data."""
     # coverage: ignore