dataclass_json_dict

cirq-core/cirq/__init__.py

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

cirq-core/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-core/cirq/protocols/__init__.py

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

cirq-core/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-core/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: