dataclass_json_dict (quantumlib#4391)

   Dataclasses keep track of their relevant fields, so we can automatically generate these.
    Dataclasses are implemented with somewhat complex metaprogramming, and tooling (PyCharm, mypy)
    have special cases for dealing with classes decorated with `@dataclass`. There is very little
    support (and no plans for support) for decorators that wrap `@dataclass` (like
    `@cirq.json_serializable_dataclass`) or combining additional decorators with `@dataclass`.
    Although not as elegant, you may want to consider explicitly defining `_json_dict_` on your
    dataclasses which simply `return dataclass_json_dict(self)`.

cirq/__init__.py

@@ -512,6 +512,7 @@
     is_parameterized,
     JsonResolver,
     json_serializable_dataclass,
+    dataclass_json_dict,
     kraus,
     measurement_key,
     measurement_key_name,

cirq/experiments/xeb_fitting.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Estimation of fidelity associated with experimental circuit executions."""
+import dataclasses
 from abc import abstractmethod, ABC
-from dataclasses import dataclass
 from typing import (
     List,
     Optional,
@@ -29,20 +29,14 @@
 import scipy.optimize
 import scipy.stats
 import sympy
-
-from cirq import ops
+from cirq import ops, protocols
 from cirq.circuits import Circuit
 from cirq.experiments.xeb_simulation import simulate_2q_xeb_circuits
 
 if TYPE_CHECKING:
     import cirq
     import multiprocessing
 
-    # Workaround for mypy custom dataclasses (python/mypy#5406)
-    from dataclasses import dataclass as json_serializable_dataclass
-else:
-    from cirq.protocols import json_serializable_dataclass
-
 THETA_SYMBOL, ZETA_SYMBOL, CHI_SYMBOL, GAMMA_SYMBOL, PHI_SYMBOL = sympy.symbols(
     'theta zeta chi gamma phi'
 )
@@ -191,8 +185,7 @@ def phased_fsim_angles_from_gate(gate: 'cirq.Gate') -> Dict[str, float]:
     raise ValueError(f"Unknown default angles for {gate}.")
 
 
-# mypy issue: https://github.com/python/mypy/issues/5374
-@json_serializable_dataclass(frozen=True)  # type: ignore
+@dataclasses.dataclass(frozen=True)
 class XEBPhasedFSimCharacterizationOptions(XEBCharacterizationOptions):
     """Options for calibrating a PhasedFSim-like gate using XEB.
 
@@ -320,6 +313,9 @@ def with_defaults_from_gate(
             **gate_to_angles_func(gate),
         )
 
+    def _json_dict_(self):
+        return protocols.dataclass_json_dict(self)
+
 
 def SqrtISwapXEBOptions(*args, **kwargs):
     """Options for calibrating a sqrt(ISWAP) gate using XEB."""
@@ -346,7 +342,7 @@ def parameterize_circuit(
 QPair_T = Tuple['cirq.Qid', 'cirq.Qid']
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class XEBCharacterizationResult:
     """The result of `characterize_phased_fsim_parameters_with_xeb`.
 
@@ -437,7 +433,7 @@ def _mean_infidelity(angles):
     )
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class _CharacterizePhasedFsimParametersWithXebClosure:
     """A closure object to wrap `characterize_phased_fsim_parameters_with_xeb` for use in
     multiprocessing."""

cirq/protocols/__init__.py

@@ -88,6 +88,7 @@
     to_json,
     read_json,
     obj_to_dict_helper,
+    dataclass_json_dict,
     SerializableByKey,
     SupportsJSON,
 )

cirq/protocols/json_serialization.py

@@ -177,6 +177,13 @@ def json_serializable_dataclass(
     the ``_json_dict_`` protocol method which automatically determines
     the appropriate fields from the dataclass.
 
+    Dataclasses are implemented with somewhat complex metaprogramming, and
+    tooling (PyCharm, mypy) have special cases for dealing with classes
+    decorated with @dataclass. There is very little support (and no plans for
+    support) for decorators that wrap @dataclass like this. Consider explicitly
+    defining `_json_dict_` on your dataclasses which simply
+    `return dataclass_json_dict(self)`.
+
     Args:
         namespace: An optional prefix to the value associated with the
             key "cirq_type". The namespace name will be joined with the
@@ -209,6 +216,21 @@ def wrap(cls):
 # pylint: enable=redefined-builtin
 
 
+def dataclass_json_dict(obj: Any, namespace: str = None) -> Dict[str, Any]:
+    """Return a dictionary suitable for _json_dict_ from a dataclass.
+
+    Dataclasses keep track of their relevant fields, so we can automatically generate these.
+
+    Dataclasses are implemented with somewhat complex metaprogramming, and tooling (PyCharm, mypy)
+    have special cases for dealing with classes decorated with @dataclass. There is very little
+    support (and no plans for support) for decorators that wrap @dataclass (like
+    @cirq.json_serializable_dataclass) or combining additional decorators with @dataclass.
+    Although not as elegant, you may want to consider explicitly defining `_json_dict_` on your
+    dataclasses which simply `return dataclass_json_dict(self)`.
+    """
+    return obj_to_dict_helper(obj, [f.name for f in dataclasses.fields(obj)], namespace=namespace)
+
+
 class CirqEncoder(json.JSONEncoder):
     """Extend json.JSONEncoder to support Cirq objects.
 

cirq/protocols/json_serialization_test.py

@@ -775,6 +775,24 @@ def custom_resolver(name):
     assert_json_roundtrip_works(my_dc, resolvers=[custom_resolver] + cirq.DEFAULT_RESOLVERS)
 
 
+def test_dataclass_json_dict():
+    @dataclasses.dataclass(frozen=True)
+    class MyDC:
+        q: cirq.LineQubit
+        desc: str
+
+        def _json_dict_(self):
+            return cirq.dataclass_json_dict(self)
+
+    def custom_resolver(name):
+        if name == 'MyDC':
+            return MyDC
+
+    my_dc = MyDC(cirq.LineQubit(4), 'hi mom')
+
+    assert_json_roundtrip_works(my_dc, resolvers=[custom_resolver, *cirq.DEFAULT_RESOLVERS])
+
+
 def test_json_serializable_dataclass_namespace():
     @cirq.json_serializable_dataclass(namespace='cirq.experiments')
     class QuantumVolumeParams: