Add return-type to public functions, mostly tests part 4

cirq-core/cirq/ops/measure_util.py

@@ -14,12 +14,12 @@
 
 from __future__ import annotations
 
-from typing import Callable, Iterable, overload, TYPE_CHECKING
+from typing import Callable, cast, Iterable, overload, TYPE_CHECKING
 
 import numpy as np
 
 from cirq import protocols
-from cirq.ops import pauli_string, raw_types
+from cirq.ops import gate_operation, pauli_string, raw_types
 from cirq.ops.measurement_gate import MeasurementGate
 from cirq.ops.pauli_measurement_gate import PauliMeasurementGate
 
@@ -96,7 +96,7 @@ def measure(
     key: str | cirq.MeasurementKey | None = None,
     invert_mask: tuple[bool, ...] = (),
     confusion_map: dict[tuple[int, ...], np.ndarray] | None = None,
-) -> raw_types.Operation:
+) -> gate_operation.GateOperation:
     pass
 
 
@@ -107,7 +107,7 @@ def measure(
     key: str | cirq.MeasurementKey | None = None,
     invert_mask: tuple[bool, ...] = (),
     confusion_map: dict[tuple[int, ...], np.ndarray] | None = None,
-) -> raw_types.Operation:
+) -> gate_operation.GateOperation:
     pass
 
 
@@ -116,7 +116,7 @@ def measure(
     key: str | cirq.MeasurementKey | None = None,
     invert_mask: tuple[bool, ...] = (),
     confusion_map: dict[tuple[int, ...], np.ndarray] | None = None,
-) -> raw_types.Operation:
+) -> gate_operation.GateOperation:
     """Returns a single MeasurementGate applied to all the given qubits.
 
     The qubits are measured in the computational basis. This can also be
@@ -161,7 +161,8 @@ def measure(
     if key is None:
         key = _default_measurement_key(targets)
     qid_shape = protocols.qid_shape(targets)
-    return MeasurementGate(len(targets), key, invert_mask, qid_shape, confusion_map).on(*targets)
+    gate = MeasurementGate(len(targets), key, invert_mask, qid_shape, confusion_map)
+    return cast(gate_operation.GateOperation, gate.on(*targets))
 
 
 M = measure

cirq-core/cirq/ops/pauli_string_test.py

@@ -2099,7 +2099,7 @@ def test_pauli_ops_identity_gate_operation(gate1: cirq.Pauli, gate2: cirq.Pauli)
     assert np.array_equal(subtraction.matrix(), unitary1 - unitary2)
 
 
-def test_pauli_gate_multiplication_with_power():
+def test_pauli_gate_multiplication_with_power() -> None:
     q = cirq.LineQubit(0)
 
     # Test all Pauli gates (X, Y, Z)
@@ -2122,7 +2122,7 @@ def test_pauli_gate_multiplication_with_power():
         assert gate**5 * gate**0 == gate**5
 
 
-def test_try_interpret_as_pauli_string():
+def test_try_interpret_as_pauli_string() -> None:
     from cirq.ops.pauli_string import _try_interpret_as_pauli_string
 
     q = cirq.LineQubit(0)

cirq-core/cirq/protocols/has_stabilizer_effect_protocol_test.py

@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import numpy as np
 
 import cirq
@@ -50,40 +52,40 @@ def __init__(self, q: cirq.Qid = cirq.LineQubit(0)):
         self.q = q
 
     @property
-    def qubits(self):
+    def qubits(self) -> tuple[cirq.Qid, ...]:
         return (self.q,)
 
-    def with_qubits(self, *new_qubits):  # pragma: no cover
+    def with_qubits(self, *new_qubits) -> cirq.Operation:  # pragma: no cover
         return self
 
 
 class NoOp(EmptyOp):
     @property
-    def gate(self):
+    def gate(self) -> Any:
         return No()
 
 
 class NoOp1(EmptyOp):
     @property
-    def gate(self):
+    def gate(self) -> Any:
         return No1()
 
 
 class NoOp2(EmptyOp):
     @property
-    def gate(self):
+    def gate(self) -> Any:
         return No2()
 
 
 class NoOp3(EmptyOp):
     @property
-    def gate(self):
+    def gate(self) -> Any:
         return No3()
 
 
 class YesOp(EmptyOp):
     @property
-    def gate(self):
+    def gate(self) -> Any:
         return Yes()
 
 
@@ -95,8 +97,8 @@ def _unitary_(self):
         return self.unitary
 
     @property
-    def qubits(self):
-        return cirq.LineQubit.range(self.unitary.shape[0].bit_length() - 1)
+    def qubits(self) -> tuple[cirq.Qid, ...]:
+        return tuple(cirq.LineQubit.range(self.unitary.shape[0].bit_length() - 1))
 
 
 class GateDecomposes(cirq.Gate):

cirq-core/cirq/protocols/has_unitary_protocol_test.py

@@ -218,8 +218,8 @@ class EmptyOp(cirq.Operation):
     """A trivial operation that will be recognized as `_apply_unitary_`-able."""
 
     @property
-    def qubits(self):
+    def qubits(self) -> tuple[cirq.Qid, ...]:
         return ()
 
-    def with_qubits(self, *new_qubits):  # pragma: no cover
-        return self
+    def with_qubits(self, *new_qubits) -> cirq.Operation:
+        raise NotImplementedError()

cirq-core/cirq/protocols/json_serialization.py

@@ -217,7 +217,7 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._memo: dict[Any, dict] = {}
 
-    def default(self, o):
+    def default(self, o) -> Any:
         # Object with custom method?
         if hasattr(o, '_json_dict_'):
             json_dict = _json_dict_with_cirq_type(o)
@@ -519,7 +519,7 @@ def read_json(
     *,
     json_text: str | None = None,
     resolvers: Sequence[JsonResolver] | None = None,
-):
+) -> Any:
     """Read a JSON file that optionally contains cirq objects.
 
     Args:
@@ -605,7 +605,7 @@ def read_json_gzip(
     *,
     gzip_raw: bytes | None = None,
     resolvers: Sequence[JsonResolver] | None = None,
-):
+) -> Any:
     """Read a gzipped JSON file that optionally contains cirq objects.
 
     Args: