diff --git a/cirq-core/cirq/__init__.py b/cirq-core/cirq/__init__.py index 9f29a232a80..a49246e83be 100644 --- a/cirq-core/cirq/__init__.py +++ b/cirq-core/cirq/__init__.py @@ -512,6 +512,7 @@ is_parameterized, JsonResolver, json_serializable_dataclass, + dataclass_json_dict, kraus, measurement_key, measurement_key_name, diff --git a/cirq-core/cirq/experiments/xeb_fitting.py b/cirq-core/cirq/experiments/xeb_fitting.py index d9e62147912..aa0860ba131 100644 --- a/cirq-core/cirq/experiments/xeb_fitting.py +++ b/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,8 +29,7 @@ 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 @@ -38,11 +37,6 @@ 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.""" diff --git a/cirq-core/cirq/protocols/__init__.py b/cirq-core/cirq/protocols/__init__.py index c9ac934d582..a9cb6e2f517 100644 --- a/cirq-core/cirq/protocols/__init__.py +++ b/cirq-core/cirq/protocols/__init__.py @@ -88,6 +88,7 @@ to_json, read_json, obj_to_dict_helper, + dataclass_json_dict, SerializableByKey, SupportsJSON, ) diff --git a/cirq-core/cirq/protocols/json_serialization.py b/cirq-core/cirq/protocols/json_serialization.py index 690d01da204..e117e3635d3 100644 --- a/cirq-core/cirq/protocols/json_serialization.py +++ b/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. diff --git a/cirq-core/cirq/protocols/json_serialization_test.py b/cirq-core/cirq/protocols/json_serialization_test.py index 897bd78adca..b47bd122630 100644 --- a/cirq-core/cirq/protocols/json_serialization_test.py +++ b/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: