Define moment+moment operation to return a combined moment (#2879)

cirq/circuits/circuit.py

@@ -210,7 +210,7 @@ def _decompose_(self) -> 'cirq.OP_TREE':
 
     # pylint: disable=function-redefined
     @overload
-    def __getitem__(self, key: slice) -> 'Circuit':
+    def __getitem__(self, key: slice) -> 'cirq.Circuit':
         pass
 
     @overload

cirq/ops/moment.py

@@ -89,7 +89,7 @@ def operates_on(self, qubits: Iterable[raw_types.Qid]) -> bool:
         """
         return bool(set(qubits) & self.qubits)
 
-    def with_operation(self, operation: 'cirq.Operation'):
+    def with_operation(self, operation: 'cirq.Operation') -> 'cirq.Moment':
         """Returns an equal moment, but with the given op added.
 
         Args:
@@ -108,7 +108,8 @@ def with_operation(self, operation: 'cirq.Operation'):
 
         return m
 
-    def without_operations_touching(self, qubits: Iterable[raw_types.Qid]):
+    def without_operations_touching(self, qubits: Iterable['cirq.Qid']
+                                   ) -> 'cirq.Moment':
         """Returns an equal moment, but without ops on the given qubits.
 
         Args:
@@ -194,8 +195,8 @@ def __str__(self):
         return ' and '.join(str(op) for op in self.operations)
 
     def transform_qubits(self: TSelf_Moment,
-                         func: Callable[[raw_types.Qid], raw_types.Qid]
-                         ) -> TSelf_Moment:
+                         func: Callable[['cirq.Qid'], 'cirq.Qid']
+                        ) -> TSelf_Moment:
         """Returns the same moment, but with different qubits.
 
         Args:
@@ -212,9 +213,12 @@ def transform_qubits(self: TSelf_Moment,
     def _json_dict_(self):
         return protocols.obj_to_dict_helper(self, ['operations'])
 
-    def __add__(self, other):
+    def __add__(self,
+                other: Union['cirq.Operation', 'cirq.Moment']) -> 'cirq.Moment':
         if isinstance(other, raw_types.Operation):
             return self.with_operation(other)
+        if isinstance(other, Moment):
+            return Moment(self.operations + other.operations)
         return NotImplemented
 
     # pylint: disable=function-redefined

cirq/ops/moment_test.py

@@ -264,7 +264,7 @@ def test_immutable_moment():
 
 
 def test_add():
-    a, b = cirq.LineQubit.range(2)
+    a, b, c = cirq.LineQubit.range(3)
     expected_circuit = cirq.Circuit([cirq.CNOT(a, b), cirq.X(a), cirq.Y(b)])
 
     circuit1 = cirq.Circuit([cirq.CNOT(a, b), cirq.X(a)])
@@ -275,6 +275,14 @@ def test_add():
     circuit2[1] += cirq.X(a)
     assert circuit2 == expected_circuit
 
+    m1 = cirq.Moment([cirq.X(a)])
+    m2 = cirq.Moment([cirq.CNOT(a, b)])
+    m3 = cirq.Moment([cirq.X(c)])
+    assert m1 + m3 == cirq.Moment([cirq.X(a), cirq.X(c)])
+    assert m2 + m3 == cirq.Moment([cirq.CNOT(a, b), cirq.X(c)])
+    with pytest.raises(ValueError, match='Overlap'):
+        _ = m1 + m2
+
 
 def test_indexes_by_qubit():
     a, b, c = cirq.LineQubit.range(3)