Make with_measurement_key_mapping live up to its name (#5610)

95-martin-orion · web-flow · commit 8f7ced71fd97 · 2022-06-24T17:09:13.000-07:00
Fixes #5552. Since this change only* affects type annotations, it is non-breaking. Any breakages due to `ParamResolver.param_dict` becoming immutable should instead refer to #5548, where this change was applied. *Okay, it also removes some `dict` calls, but those were added in #5548.
diff --git a/cirq-core/cirq/circuits/circuit.py b/cirq-core/cirq/circuits/circuit.py
@@ -29,6 +29,7 @@
     AbstractSet,
     Any,
     Callable,
+    Mapping,
     cast,
     Dict,
     FrozenSet,
@@ -944,7 +945,7 @@ def all_measurement_key_names(self) -> FrozenSet[str]:
     def _measurement_key_names_(self) -> FrozenSet[str]:
         return self.all_measurement_key_names()
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         return self._with_sliced_moments(
             [protocols.with_measurement_key_mapping(moment, key_map) for moment in self.moments]
         )
diff --git a/cirq-core/cirq/circuits/circuit_operation.py b/cirq-core/cirq/circuits/circuit_operation.py
@@ -322,7 +322,7 @@ def _measurement_key_objs(self) -> FrozenSet['cirq.MeasurementKey']:
             key.with_key_path_prefix(*self.parent_path) for key in circuit_keys
         )
         return frozenset(
-            protocols.with_measurement_key_mapping(key, dict(self.measurement_key_map))
+            protocols.with_measurement_key_mapping(key, self.measurement_key_map)
             for key in circuit_keys
         )
 
@@ -368,9 +368,7 @@ def _mapped_any_loop(self) -> 'cirq.Circuit':
         if isinstance(self.repetitions, INT_CLASSES) and self.repetitions < 0:
             circuit = circuit**-1
         if self.measurement_key_map:
-            circuit = protocols.with_measurement_key_mapping(
-                circuit, dict(self.measurement_key_map)
-            )
+            circuit = protocols.with_measurement_key_mapping(circuit, self.measurement_key_map)
         if self.param_resolver:
             circuit = protocols.resolve_parameters(circuit, self.param_resolver, recursive=False)
         return circuit.unfreeze(copy=False)
diff --git a/cirq-core/cirq/circuits/moment.py b/cirq-core/cirq/circuits/moment.py
@@ -22,6 +22,7 @@
     FrozenSet,
     Iterable,
     Iterator,
+    Mapping,
     overload,
     Optional,
     Sequence,
@@ -229,7 +230,7 @@ def without_operations_touching(self, qubits: Iterable['cirq.Qid']) -> 'cirq.Mom
             if qubits.isdisjoint(frozenset(operation.qubits))
         )
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         return Moment(
             protocols.with_measurement_key_mapping(op, key_map)
             if protocols.measurement_keys_touched(op)
diff --git a/cirq-core/cirq/ops/classically_controlled_operation.py b/cirq-core/cirq/ops/classically_controlled_operation.py
@@ -14,6 +14,7 @@
 from typing import (
     AbstractSet,
     Any,
+    Mapping,
     cast,
     Dict,
     FrozenSet,
@@ -178,7 +179,7 @@ def _act_on_(self, sim_state: 'cirq.SimulationStateBase') -> bool:
         return True
 
     def _with_measurement_key_mapping_(
-        self, key_map: Dict[str, str]
+        self, key_map: Mapping[str, str]
     ) -> 'ClassicallyControlledOperation':
         conditions = [protocols.with_measurement_key_mapping(c, key_map) for c in self._conditions]
         sub_operation = protocols.with_measurement_key_mapping(self._sub_operation, key_map)
diff --git a/cirq-core/cirq/ops/gate_operation.py b/cirq-core/cirq/ops/gate_operation.py
@@ -18,6 +18,7 @@
 from typing import (
     AbstractSet,
     Any,
+    Mapping,
     cast,
     Collection,
     Dict,
@@ -81,7 +82,7 @@ def with_gate(self, new_gate: 'cirq.Gate') -> 'cirq.Operation':
             return self
         return new_gate.on(*self.qubits)
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         new_gate = protocols.with_measurement_key_mapping(self.gate, key_map)
         if new_gate is NotImplemented:
             return NotImplemented
diff --git a/cirq-core/cirq/ops/kraus_channel.py b/cirq-core/cirq/ops/kraus_channel.py
@@ -1,5 +1,5 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
-from typing import Any, Dict, FrozenSet, Iterable, Tuple, TYPE_CHECKING, Union
+from typing import Any, Dict, FrozenSet, Iterable, Mapping, Tuple, TYPE_CHECKING, Union
 import numpy as np
 
 from cirq import linalg, protocols, value
@@ -84,7 +84,7 @@ def _measurement_key_obj_(self) -> 'cirq.MeasurementKey':
             return NotImplemented
         return self._key
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         if self._key is None:
             return NotImplemented
         if self._key not in key_map:
diff --git a/cirq-core/cirq/ops/measurement_gate.py b/cirq-core/cirq/ops/measurement_gate.py
@@ -12,7 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, FrozenSet, Iterable, Optional, Tuple, Sequence, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Dict,
+    FrozenSet,
+    Iterable,
+    Mapping,
+    Optional,
+    Tuple,
+    Sequence,
+    TYPE_CHECKING,
+    Union,
+)
 
 import numpy as np
 
@@ -129,7 +140,7 @@ def _with_rescoped_keys_(
     ):
         return self.with_key(protocols.with_rescoped_keys(self.mkey, path, bindable_keys))
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         return self.with_key(protocols.with_measurement_key_mapping(self.mkey, key_map))
 
     def with_bits_flipped(self, *bit_positions: int) -> 'MeasurementGate':
diff --git a/cirq-core/cirq/ops/mixed_unitary_channel.py b/cirq-core/cirq/ops/mixed_unitary_channel.py
@@ -1,5 +1,5 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
-from typing import Any, Dict, FrozenSet, Iterable, Tuple, TYPE_CHECKING, Union
+from typing import Any, Dict, FrozenSet, Iterable, Mapping, Tuple, TYPE_CHECKING, Union
 import numpy as np
 
 from cirq import linalg, protocols, value
@@ -87,7 +87,7 @@ def _measurement_key_obj_(self) -> 'cirq.MeasurementKey':
             return NotImplemented
         return self._key
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         if self._key is None:
             return NotImplemented
         if self._key not in key_map:
diff --git a/cirq-core/cirq/ops/pauli_measurement_gate.py b/cirq-core/cirq/ops/pauli_measurement_gate.py
@@ -12,7 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, FrozenSet, Iterable, Tuple, Sequence, TYPE_CHECKING, Union, cast
+from typing import (
+    Any,
+    Dict,
+    FrozenSet,
+    Iterable,
+    Mapping,
+    Tuple,
+    Sequence,
+    TYPE_CHECKING,
+    Union,
+    cast,
+)
 
 from cirq import protocols, value
 from cirq.ops import (
@@ -103,7 +114,7 @@ def _with_rescoped_keys_(
     ) -> 'PauliMeasurementGate':
         return self.with_key(protocols.with_rescoped_keys(self.mkey, path, bindable_keys))
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]) -> 'PauliMeasurementGate':
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]) -> 'PauliMeasurementGate':
         return self.with_key(protocols.with_measurement_key_mapping(self.mkey, key_map))
 
     def with_observable(
diff --git a/cirq-core/cirq/ops/raw_types.py b/cirq-core/cirq/ops/raw_types.py
@@ -26,6 +26,7 @@
     Hashable,
     Iterable,
     List,
+    Mapping,
     Optional,
     Sequence,
     Tuple,
@@ -733,7 +734,7 @@ def gate(self) -> Optional['cirq.Gate']:
     def with_qubits(self, *new_qubits: 'cirq.Qid'):
         return TaggedOperation(self.sub_operation.with_qubits(*new_qubits), *self._tags)
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         sub_op = protocols.with_measurement_key_mapping(self.sub_operation, key_map)
         if sub_op is NotImplemented:
             return NotImplemented
diff --git a/cirq-core/cirq/protocols/measurement_key_protocol.py b/cirq-core/cirq/protocols/measurement_key_protocol.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Protocol for object that have measurement keys."""
 
-from typing import Any, Dict, FrozenSet, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, FrozenSet, Mapping, Optional, Tuple, TYPE_CHECKING, Union
 
 from typing_extensions import Protocol
 
@@ -98,7 +98,7 @@ def _measurement_key_names_(self) -> Union[FrozenSet[str], NotImplementedType, N
         """
 
     @doc_private
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         """Return a copy of this object with the measurement keys remapped.
 
         This method allows measurement keys to be reassigned at runtime.
@@ -291,7 +291,7 @@ def is_measurement(val: Any) -> bool:
     return keys is not NotImplemented and bool(keys)
 
 
-def with_measurement_key_mapping(val: Any, key_map: Dict[str, str]):
+def with_measurement_key_mapping(val: Any, key_map: Mapping[str, str]):
     """Remaps the target's measurement keys according to the provided key_map.
 
     This method can be used to reassign measurement keys at runtime, or to
diff --git a/cirq-core/cirq/value/condition.py b/cirq-core/cirq/value/condition.py
@@ -14,7 +14,7 @@
 
 import abc
 import dataclasses
-from typing import Dict, Tuple, TYPE_CHECKING, FrozenSet
+from typing import Mapping, Tuple, TYPE_CHECKING, FrozenSet
 
 import sympy
 
@@ -47,7 +47,7 @@ def resolve(self, classical_data: 'cirq.ClassicalDataStoreReader') -> bool:
     def qasm(self):
         """Returns the qasm of this condition."""
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]) -> 'cirq.Condition':
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]) -> 'cirq.Condition':
         condition = self
         for k in self.keys:
             condition = condition.replace_key(k, mkp.with_measurement_key_mapping(k, key_map))
diff --git a/cirq-core/cirq/value/measurement_key.py b/cirq-core/cirq/value/measurement_key.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, FrozenSet, Optional, Tuple
+from typing import FrozenSet, Mapping, Optional, Tuple
 
 import dataclasses
 
@@ -122,7 +122,7 @@ def _with_rescoped_keys_(
     ):
         return self.replace(path=path + self.path)
 
-    def _with_measurement_key_mapping_(self, key_map: Dict[str, str]):
+    def _with_measurement_key_mapping_(self, key_map: Mapping[str, str]):
         if self.name not in key_map:
             return self
         return self.replace(name=key_map[self.name])
