cirq-core 1.4.1__py3-none-any.whl → 1.5.0__py3-none-any.whl

cirq/ops/identity.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 """IdentityGate."""
 
-from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Sequence
+import numbers
+from types import NotImplementedType
+from typing import Any, Dict, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 import sympy
@@ -71,10 +73,16 @@ class IdentityGate(raw_types.Gate):
         return len(self._qid_shape)
 
     def __pow__(self, power: Any) -> Any:
-        if isinstance(power, (int, float, complex, sympy.Basic)):
+        if isinstance(power, (numbers.Complex, sympy.Basic)):
             return self
         return NotImplemented
 
+    def _commutes_(self, other: Any, *, atol: float = 1e-8) -> Union[bool, NotImplementedType]:
+        """The identity gate commutes with all other gates."""
+        if not isinstance(other, raw_types.Gate):
+            return NotImplemented
+        return True
+
     def _has_unitary_(self) -> bool:
         return True
 
@@ -119,7 +127,7 @@ class IdentityGate(raw_types.Gate):
     def _mul_with_qubits(self, qubits: Tuple['cirq.Qid', ...], other):
         if isinstance(other, raw_types.Operation):
             return other
-        if isinstance(other, (complex, float, int)):
+        if isinstance(other, numbers.Complex):
             from cirq.ops.pauli_string import PauliString
 
             return PauliString(coefficient=other)
@@ -128,10 +136,12 @@ class IdentityGate(raw_types.Gate):
     _rmul_with_qubits = _mul_with_qubits
 
     def _circuit_diagram_info_(self, args) -> Tuple[str, ...]:
+        if self.num_qubits() <= 0:
+            return NotImplemented
         return ('I',) * self.num_qubits()
 
     def _qasm_(self, args: 'cirq.QasmArgs', qubits: Tuple['cirq.Qid', ...]) -> Optional[str]:
-        args.validate_version('2.0')
+        args.validate_version('2.0', '3.0')
         return ''.join([args.format('id {0};\n', qubit) for qubit in qubits])
 
     @classmethod

cirq/ops/identity_test.py

@@ -208,3 +208,27 @@ def test_identity_short_circuits_act_on():
     args = mock.Mock(cirq.SimulationState)
     args._act_on_fallback_.side_effect = mock.Mock(side_effect=Exception('No!'))
     cirq.act_on(cirq.IdentityGate(1)(cirq.LineQubit(0)), args)
+
+
+def test_identity_commutes():
+    assert cirq.commutes(cirq.I, cirq.X)
+    with pytest.raises(TypeError):
+        cirq.commutes(cirq.I, "Gate")
+
+
+def test_identity_diagram():
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(cirq.IdentityGate(3).on_each(cirq.LineQubit.range(3))),
+        """
+0: ───I───
+      │
+1: ───I───
+      │
+2: ───I───
+""",
+    )
+    cirq.testing.assert_has_diagram(
+        cirq.Circuit(cirq.IdentityGate(0)()),
+        """
+    I(0)""",
+    )

cirq/ops/kraus_channel.py

@@ -1,5 +1,6 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
 from typing import Any, Dict, FrozenSet, Iterable, Mapping, Tuple, TYPE_CHECKING, Union
+
 import numpy as np
 
 from cirq import linalg, protocols, value

cirq/ops/kraus_channel_test.py

@@ -1,8 +1,9 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
-import cirq
 import numpy as np
 import pytest
 
+import cirq
+
 
 def test_kraus_channel_from_channel():
     q0 = cirq.LineQubit(0)
@@ -53,6 +54,7 @@ def test_kraus_channel_remap_keys():
     kc_x = cirq.KrausChannel.from_channel(dp, key='x')
     assert cirq.with_measurement_key_mapping(kc_x, {'a': 'b'}) is kc_x
     assert cirq.measurement_key_name(cirq.with_key_path(kc_x, ('path',))) == 'path:x'
+    assert cirq.measurement_key_name(cirq.with_key_path_prefix(kc_x, ('path',))) == 'path:x'
 
     kc_a = cirq.KrausChannel.from_channel(dp, key='a')
     kc_b = cirq.KrausChannel.from_channel(dp, key='b')

cirq/ops/linear_combinations.py

@@ -12,34 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numbers
 from collections import defaultdict
 from typing import (
     AbstractSet,
     Any,
+    DefaultDict,
     Dict,
+    FrozenSet,
     Iterable,
+    List,
     Mapping,
     Optional,
     Tuple,
-    Union,
-    List,
-    FrozenSet,
-    DefaultDict,
     TYPE_CHECKING,
+    Union,
 )
-import numbers
 
 import numpy as np
-from sympy.logic.boolalg import And, Not, Or, Xor
+from scipy.sparse import csr_matrix
 from sympy.core.expr import Expr
 from sympy.core.symbol import Symbol
-from scipy.sparse import csr_matrix
+from sympy.logic.boolalg import And, Not, Or, Xor
 
 from cirq import linalg, protocols, qis, value
 from cirq._doc import document
 from cirq.linalg import operator_spaces
-from cirq.ops import identity, raw_types, pauli_gates, pauli_string
-from cirq.ops.pauli_string import PauliString, _validate_qubit_mapping
+from cirq.ops import identity, pauli_gates, pauli_string, raw_types
+from cirq.ops.pauli_string import _validate_qubit_mapping, PauliString
 from cirq.ops.projector import ProjectorString
 from cirq.value.linear_dict import _format_terms
 
@@ -48,7 +48,7 @@ if TYPE_CHECKING:
 
 UnitPauliStringT = FrozenSet[Tuple[raw_types.Qid, pauli_gates.Pauli]]
 PauliSumLike = Union[
-    int, float, complex, PauliString, 'PauliSum', pauli_string.SingleQubitPauliStringGateOperation
+    complex, PauliString, 'PauliSum', pauli_string.SingleQubitPauliStringGateOperation
 ]
 document(
     PauliSumLike,
@@ -81,7 +81,7 @@ class LinearCombinationOfGates(value.LinearDict[raw_types.Gate]):
         2 * cirq.X - 2 * cirq.Z
     """
 
-    def __init__(self, terms: Mapping[raw_types.Gate, value.Scalar]) -> None:
+    def __init__(self, terms: Mapping[raw_types.Gate, 'cirq.TParamValComplex']) -> None:
         """Initializes linear combination from a collection of terms.
 
         Args:
@@ -149,17 +149,19 @@ class LinearCombinationOfGates(value.LinearDict[raw_types.Gate]):
         )
 
     def _is_parameterized_(self) -> bool:
-        return any(protocols.is_parameterized(gate) for gate in self.keys())
+        return any(protocols.is_parameterized(item) for item in self.items())
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return {name for gate in self.keys() for name in protocols.parameter_names(gate)}
+        return {name for item in self.items() for name in protocols.parameter_names(item)}
 
     def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'LinearCombinationOfGates':
         return self.__class__(
             {
-                protocols.resolve_parameters(gate, resolver, recursive): coeff
+                protocols.resolve_parameters(
+                    gate, resolver, recursive
+                ): protocols.resolve_parameters(coeff, resolver, recursive)
                 for gate, coeff in self.items()
             }
         )
@@ -222,7 +224,7 @@ class LinearCombinationOfOperations(value.LinearDict[raw_types.Operation]):
     by the identity operator. Note that A may not be unitary or even normal.
     """
 
-    def __init__(self, terms: Mapping[raw_types.Operation, value.Scalar]) -> None:
+    def __init__(self, terms: Mapping[raw_types.Operation, 'cirq.TParamValComplex']) -> None:
         """Initializes linear combination from a collection of terms.
 
         Args:
@@ -264,17 +266,19 @@ class LinearCombinationOfOperations(value.LinearDict[raw_types.Operation]):
         return LinearCombinationOfOperations({i: bi, x: bx, y: by, z: bz})
 
     def _is_parameterized_(self) -> bool:
-        return any(protocols.is_parameterized(op) for op in self.keys())
+        return any(protocols.is_parameterized(item) for item in self.items())
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return {name for op in self.keys() for name in protocols.parameter_names(op)}
+        return {name for item in self.items() for name in protocols.parameter_names(item)}
 
     def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'LinearCombinationOfOperations':
         return self.__class__(
             {
-                protocols.resolve_parameters(op, resolver, recursive): coeff
+                protocols.resolve_parameters(op, resolver, recursive): protocols.resolve_parameters(
+                    coeff, resolver, recursive
+                )
                 for op, coeff in self.items()
             }
         )
@@ -353,7 +357,9 @@ def _is_linear_dict_of_unit_pauli_string(linear_dict: value.LinearDict[UnitPauli
     return True
 
 
-def _pauli_string_from_unit(unit: UnitPauliStringT, coefficient: Union[int, float, complex] = 1):
+def _pauli_string_from_unit(
+    unit: UnitPauliStringT, coefficient: Union[int, float, 'cirq.TParamValComplex'] = 1
+):
     return PauliString(qubit_pauli_map=dict(unit), coefficient=coefficient)
 
 
@@ -805,7 +811,7 @@ class PauliSum:
         if not isinstance(other, (numbers.Complex, PauliString, PauliSum)):
             return NotImplemented
         if isinstance(other, numbers.Complex):
-            self._linear_dict *= other
+            self._linear_dict *= complex(other)
         elif isinstance(other, PauliString):
             temp = PauliSum.from_pauli_strings([term * other for term in self])
             self._linear_dict = temp._linear_dict
@@ -870,7 +876,7 @@ def _projector_string_from_projector_dict(projector_dict, coefficient=1.0):
     return ProjectorString(dict(projector_dict), coefficient)
 
 
-@value.value_equality(approximate=True)
+@value.value_equality(approximate=True, unhashable=True)
 class ProjectorSum:
     """List of mappings representing a sum of projector operators."""
 

cirq/ops/linear_combinations_test.py

@@ -153,6 +153,10 @@ def test_non_unitary_linear_combination_of_gates_has_no_unitary(terms):
         ),
         ({cirq.X: 2, cirq.H: 1}, {'X': 2 + np.sqrt(0.5), 'Z': np.sqrt(0.5)}),
         ({cirq.XX: -2, cirq.YY: 3j, cirq.ZZ: 4}, {'XX': -2, 'YY': 3j, 'ZZ': 4}),
+        (
+            {cirq.X: sympy.Symbol('x'), cirq.Y: -sympy.Symbol('y')},
+            {'X': sympy.Symbol('x'), 'Y': -sympy.Symbol('y')},
+        ),
     ),
 )
 def test_linear_combination_of_gates_has_correct_pauli_expansion(terms, expected_expansion):
@@ -206,7 +210,11 @@ def test_linear_combinations_of_gates_invalid_powers(terms, exponent):
 
 @pytest.mark.parametrize(
     'terms, is_parameterized, parameter_names',
-    [({cirq.H: 1}, False, set()), ({cirq.X ** sympy.Symbol('t'): 1}, True, {'t'})],
+    [
+        ({cirq.H: 1}, False, set()),
+        ({cirq.X ** sympy.Symbol('t'): 1}, True, {'t'}),
+        ({cirq.X: sympy.Symbol('t')}, True, {'t'}),
+    ],
 )
 @pytest.mark.parametrize('resolve_fn', [cirq.resolve_parameters, cirq.resolve_parameters_once])
 def test_parameterized_linear_combination_of_gates(
@@ -225,7 +233,7 @@ def get_matrix(
         cirq.GateOperation,
         cirq.LinearCombinationOfGates,
         cirq.LinearCombinationOfOperations,
-    ]
+    ],
 ) -> np.ndarray:
     if isinstance(operator, (cirq.LinearCombinationOfGates, cirq.LinearCombinationOfOperations)):
         return operator.matrix()
@@ -243,13 +251,13 @@ def assert_linear_combinations_are_equal(
 
     actual_matrix = get_matrix(actual)
     expected_matrix = get_matrix(expected)
-    assert np.allclose(actual_matrix, expected_matrix)
+    assert cirq.approx_eq(actual_matrix, expected_matrix)
 
     actual_expansion = cirq.pauli_expansion(actual)
     expected_expansion = cirq.pauli_expansion(expected)
     assert set(actual_expansion.keys()) == set(expected_expansion.keys())
     for name in actual_expansion.keys():
-        assert abs(actual_expansion[name] - expected_expansion[name]) < 1e-12
+        assert cirq.approx_eq(actual_expansion[name], expected_expansion[name])
 
 
 @pytest.mark.parametrize(
@@ -279,6 +287,8 @@ def assert_linear_combinations_are_equal(
         ),
         ((cirq.X + cirq.Y + cirq.Z) ** 0, cirq.I),
         ((cirq.X - 1j * cirq.Y) ** 0, cirq.I),
+        (cirq.Y - sympy.Symbol('s') * cirq.Y, (1 - sympy.Symbol('s')) * cirq.Y),
+        ((cirq.X + cirq.Z) * sympy.Symbol('s') / np.sqrt(2), cirq.H * sympy.Symbol('s')),
     ),
 )
 def test_gate_expressions(expression, expected_result):
@@ -659,6 +669,10 @@ def test_non_unitary_linear_combination_of_operations_has_no_unitary(terms):
             {'IIZI': 1, 'IZII': 1, 'IZZI': -1},
         ),
         ({cirq.CNOT(q0, q1): 2, cirq.Z(q0): -1, cirq.X(q1): -1}, {'II': 1, 'ZX': -1}),
+        (
+            {cirq.X(q0): -sympy.Symbol('x'), cirq.Y(q0): sympy.Symbol('y')},
+            {'X': -sympy.Symbol('x'), 'Y': sympy.Symbol('y')},
+        ),
     ),
 )
 def test_linear_combination_of_operations_has_correct_pauli_expansion(terms, expected_expansion):
@@ -716,6 +730,7 @@ def test_linear_combinations_of_operations_invalid_powers(terms, exponent):
     [
         ({cirq.H(cirq.LineQubit(0)): 1}, False, set()),
         ({cirq.X(cirq.LineQubit(0)) ** sympy.Symbol('t'): 1}, True, {'t'}),
+        ({cirq.X(cirq.LineQubit(0)): sympy.Symbol('t')}, True, {'t'}),
     ],
 )
 @pytest.mark.parametrize('resolve_fn', [cirq.resolve_parameters, cirq.resolve_parameters_once])
@@ -788,6 +803,10 @@ def test_parameterized_linear_combination_of_ops(
             cirq.LinearCombinationOfOperations({cirq.X(q1): 2, cirq.Z(q1): 3}) ** 0,
             cirq.LinearCombinationOfOperations({cirq.I(q1): 1}),
         ),
+        (
+            cirq.LinearCombinationOfOperations({cirq.X(q0): sympy.Symbol('s')}) ** 2,
+            cirq.LinearCombinationOfOperations({cirq.I(q0): sympy.Symbol('s') ** 2}),
+        ),
     ),
 )
 def test_operation_expressions(expression, expected_result):

cirq/ops/matrix_gates.py

@@ -14,13 +14,13 @@
 
 """Quantum gates defined by a matrix."""
 
-from typing import Any, Dict, Iterable, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING
 
 import numpy as np
 
-from cirq import linalg, protocols, _import
+from cirq import _import, linalg, protocols
 from cirq._compat import proper_repr
-from cirq.ops import raw_types, phased_x_z_gate
+from cirq.ops import global_phase_op, identity, phased_x_z_gate, raw_types
 
 if TYPE_CHECKING:
     import cirq
@@ -148,18 +148,34 @@ class MatrixGate(raw_types.Gate):
         return MatrixGate(matrix=result.reshape(self._matrix.shape), qid_shape=self._qid_shape)
 
     def _decompose_(self, qubits: Tuple['cirq.Qid', ...]) -> 'cirq.OP_TREE':
+        from cirq.circuits import Circuit
+
+        decomposed: List['cirq.Operation'] = NotImplemented
         if self._qid_shape == (2,):
-            return [
+            decomposed = [
                 g.on(qubits[0])
                 for g in single_qubit_decompositions.single_qubit_matrix_to_gates(self._matrix)
             ]
         if self._qid_shape == (2,) * 2:
-            return two_qubit_to_cz.two_qubit_matrix_to_cz_operations(
+            decomposed = two_qubit_to_cz.two_qubit_matrix_to_cz_operations(
                 *qubits, self._matrix, allow_partial_czs=True
             )
         if self._qid_shape == (2,) * 3:
-            return three_qubit_decomposition.three_qubit_matrix_to_operations(*qubits, self._matrix)
-        return NotImplemented
+            decomposed = three_qubit_decomposition.three_qubit_matrix_to_operations(
+                *qubits, self._matrix
+            )
+        if decomposed is NotImplemented:
+            return NotImplemented
+        # The above algorithms ignore phase, but phase is important to maintain if the gate is
+        # controlled. Here, we add it back in with a global phase op.
+        ident = identity.IdentityGate(qid_shape=self._qid_shape).on(*qubits)  # Preserve qid order
+        u = protocols.unitary(Circuit(ident, *decomposed)).reshape(self._matrix.shape)
+        phase_delta = linalg.phase_delta(u, self._matrix)
+        # Phase delta is on the complex unit circle, so if real(phase_delta) >= 1, that means
+        # no phase delta. (>1 is rounding error).
+        if phase_delta.real < 1:
+            decomposed.append(global_phase_op.global_phase_operation(phase_delta))
+        return decomposed
 
     def _has_unitary_(self) -> bool:
         return True
@@ -181,7 +197,7 @@ class MatrixGate(raw_types.Gate):
         return protocols.CircuitDiagramInfo(wire_symbols=[main, *rest])
 
     def _qasm_(self, args: 'cirq.QasmArgs', qubits: Tuple['cirq.Qid', ...]) -> Optional[str]:
-        args.validate_version('2.0')
+        args.validate_version('2.0', '3.0')
         if self._qid_shape == (2,):
             return protocols.qasm(
                 phased_x_z_gate.PhasedXZGate.from_matrix(self._matrix), args=args, qubits=qubits

cirq/ops/measure_util.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Dict, Iterable, List, overload, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Callable, Dict, Iterable, List, Optional, overload, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 
 from cirq import protocols
-from cirq.ops import raw_types, pauli_string
+from cirq.ops import pauli_string, raw_types
 from cirq.ops.measurement_gate import MeasurementGate
 from cirq.ops.pauli_measurement_gate import PauliMeasurementGate
 

cirq/ops/measurement_gate.py

@@ -19,8 +19,8 @@ from typing import (
     Iterable,
     Mapping,
     Optional,
-    Tuple,
     Sequence,
+    Tuple,
     TYPE_CHECKING,
     Union,
 )
@@ -227,7 +227,7 @@ class MeasurementGate(raw_types.Gate):
     def _qasm_(self, args: 'cirq.QasmArgs', qubits: Tuple['cirq.Qid', ...]) -> Optional[str]:
         if self.confusion_map or not all(d == 2 for d in self._qid_shape):
             return NotImplemented
-        args.validate_version('2.0')
+        args.validate_version('2.0', '3.0')
         invert_mask = self.invert_mask
         if len(invert_mask) < len(qubits):
             invert_mask = invert_mask + (False,) * (len(qubits) - len(invert_mask))
@@ -235,7 +235,10 @@ class MeasurementGate(raw_types.Gate):
         for i, (qubit, inv) in enumerate(zip(qubits, invert_mask)):
             if inv:
                 lines.append(args.format('x {0};  // Invert the following measurement\n', qubit))
-            lines.append(args.format('measure {0} -> {1:meas}[{2}];\n', qubit, self.key, i))
+            if args.version == '2.0':
+                lines.append(args.format('measure {0} -> {1:meas}[{2}];\n', qubit, self.key, i))
+            else:
+                lines.append(args.format('{1:meas}[{2}] = measure {0};\n', qubit, self.key, i))
             if inv:
                 lines.append(args.format('x {0};  // Undo the inversion\n', qubit))
         return ''.join(lines)
@@ -270,7 +273,7 @@ class MeasurementGate(raw_types.Gate):
             (idxs, tuple(v for _, v in np.ndenumerate(cmap)))
             for idxs, cmap in self._confusion_map.items()
         )
-        return self.key, self.invert_mask, self._qid_shape, hashable_cmap
+        return self.key, self.full_invert_mask(), self._qid_shape, hashable_cmap
 
     def _json_dict_(self) -> Dict[str, Any]:
         other: Dict[str, Any] = {}

cirq/ops/measurement_gate_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import cast
+
 import numpy as np
 import pytest
 
@@ -74,11 +75,11 @@ def test_measurement_eq():
     eq.make_equality_group(
         lambda: cirq.MeasurementGate(1, 'a'),
         lambda: cirq.MeasurementGate(1, 'a', invert_mask=()),
+        lambda: cirq.MeasurementGate(1, 'a', invert_mask=(False,)),
         lambda: cirq.MeasurementGate(1, 'a', qid_shape=(2,)),
         lambda: cirq.MeasurementGate(1, 'a', confusion_map={}),
     )
     eq.add_equality_group(cirq.MeasurementGate(1, 'a', invert_mask=(True,)))
-    eq.add_equality_group(cirq.MeasurementGate(1, 'a', invert_mask=(False,)))
     eq.add_equality_group(
         cirq.MeasurementGate(1, 'a', confusion_map={(0,): np.array([[0, 1], [1, 0]])})
     )

cirq/ops/mixed_unitary_channel.py

@@ -1,5 +1,6 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
 from typing import Any, Dict, FrozenSet, Iterable, Mapping, Tuple, TYPE_CHECKING, Union
+
 import numpy as np
 
 from cirq import linalg, protocols, value

cirq/ops/mixed_unitary_channel_test.py

@@ -1,8 +1,9 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
-import cirq
 import numpy as np
 import pytest
 
+import cirq
+
 
 def test_matrix_mixture_from_mixture():
     q0 = cirq.LineQubit(0)
@@ -54,6 +55,7 @@ def test_matrix_mixture_remap_keys():
     mm_x = cirq.MixedUnitaryChannel.from_mixture(dp, key='x')
     assert cirq.with_measurement_key_mapping(mm_x, {'a': 'b'}) is mm_x
     assert cirq.measurement_key_name(cirq.with_key_path(mm_x, ('path',))) == 'path:x'
+    assert cirq.measurement_key_name(cirq.with_key_path_prefix(mm_x, ('path',))) == 'path:x'
 
     mm_a = cirq.MixedUnitaryChannel.from_mixture(dp, key='a')
     mm_b = cirq.MixedUnitaryChannel.from_mixture(dp, key='b')

cirq/ops/named_qubit.py

@@ -18,7 +18,6 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 from cirq import protocols
 from cirq.ops import raw_types
 
-
 if TYPE_CHECKING:
     import cirq
 
@@ -134,6 +133,10 @@ class NamedQid(_BaseNamedQid):
         """Returns a tuple of args to pass to __new__ when unpickling."""
         return (self._name, self._dimension)
 
+    # avoid pickling the _hash value, attributes are already stored with __getnewargs__
+    def __getstate__(self) -> Dict[str, Any]:
+        return {}
+
     def __repr__(self) -> str:
         return f'cirq.NamedQid({self._name!r}, dimension={self._dimension})'
 
@@ -202,6 +205,10 @@ class NamedQubit(_BaseNamedQid):
         """Returns a tuple of args to pass to __new__ when unpickling."""
         return (self._name,)
 
+    # avoid pickling the _hash value, attributes are already stored with __getnewargs__
+    def __getstate__(self) -> Dict[str, Any]:
+        return {}
+
     def __str__(self) -> str:
         return self._name
 

cirq/ops/op_tree.py

@@ -12,11 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A recursive type describing trees of operations, and utility methods for it.
-"""
+"""A recursive type describing trees of operations, and utility methods for it."""
 
-from typing import Callable, Iterable, Iterator, NoReturn, Union, TYPE_CHECKING
-from typing_extensions import Protocol
+from typing import Callable, Iterable, Iterator, NoReturn, TYPE_CHECKING, Union
 
 from cirq._doc import document
 from cirq._import import LazyLoader
@@ -28,32 +26,7 @@ if TYPE_CHECKING:
 moment = LazyLoader("moment", globals(), "cirq.circuits.moment")
 
 
-class OpTree(Protocol):
-    """The recursive type consumed by circuit builder methods.
-
-    An OpTree is a type protocol, satisfied by anything that can be recursively
-    flattened into Operations. We also define the Union type OP_TREE which
-    can be an OpTree or just a single Operation.
-
-    For example:
-    - An Operation is an OP_TREE all by itself.
-    - A list of operations is an OP_TREE.
-    - A list of tuples of operations is an OP_TREE.
-    - A list with a mix of operations and lists of operations is an OP_TREE.
-    - A generator yielding operations is an OP_TREE.
-
-    Note: once mypy supports recursive types this could be defined as an alias:
-
-    OP_TREE = Union[Operation, Iterable['OP_TREE']]
-
-    See: https://github.com/python/mypy/issues/731
-    """
-
-    def __iter__(self) -> Iterator[Union[Operation, 'OpTree']]:
-        pass
-
-
-OP_TREE = Union[Operation, OpTree]
+OP_TREE = Union[Operation, Iterable['OP_TREE']]
 document(
     OP_TREE,
     """An operation or nested collections of operations.

cirq/ops/op_tree_test.py

@@ -135,6 +135,10 @@ def test_transform_leaves():
     # Just an item.
     assert move_tree_left_freeze(operations[0]) == expected[0]
 
+    # Just a moment
+    m = cirq.Moment(cirq.X(cirq.q(1)))
+    assert cirq.transform_op_tree(m, preserve_moments=True) is m
+
     # Flat list.
     assert move_tree_left_freeze(operations) == tuple(expected)
 

cirq/ops/parallel_gate.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from typing import AbstractSet, Union, Any, Optional, Tuple, TYPE_CHECKING, Dict
+from types import NotImplementedType
+from typing import AbstractSet, Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 
 from cirq import protocols, value
 from cirq.ops import raw_types
-from cirq.type_workarounds import NotImplementedType
 
 if TYPE_CHECKING:
     import cirq

cirq/ops/parallel_gate_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
 import numpy as np
+import pytest
 import sympy
+
 import cirq