cirq-core 1.7.0.dev20250825174419__py3-none-any.whl → 1.7.0.dev20251203004401__py3-none-any.whl

cirq/protocols/json_serialization_test.py

@@ -72,7 +72,7 @@ def _get_testspecs_for_modules() -> list[ModuleJsonTestSpec]:
 MODULE_TEST_SPECS = _get_testspecs_for_modules()
 
 
-def test_deprecated_cirq_type_in_json_dict():
+def test_deprecated_cirq_type_in_json_dict() -> None:
     class HasOldJsonDict:
         # Required for testing serialization of non-cirq objects.
         __module__ = 'test.noncirq.namespace'
@@ -99,7 +99,7 @@ def test_deprecated_cirq_type_in_json_dict():
         assert_json_roundtrip_works(HasOldJsonDict(), resolvers=test_resolvers)
 
 
-def test_line_qubit_roundtrip():
+def test_line_qubit_roundtrip() -> None:
     q1 = cirq.LineQubit(12)
     assert_json_roundtrip_works(
         q1,
@@ -110,7 +110,7 @@ def test_line_qubit_roundtrip():
     )
 
 
-def test_gridqubit_roundtrip():
+def test_gridqubit_roundtrip() -> None:
     q = cirq.GridQubit(15, 18)
     assert_json_roundtrip_works(
         q,
@@ -122,7 +122,7 @@ def test_gridqubit_roundtrip():
     )
 
 
-def test_op_roundtrip():
+def test_op_roundtrip() -> None:
     q = cirq.LineQubit(5)
     op1 = cirq.rx(0.123).on(q)
     assert_json_roundtrip_works(
@@ -143,7 +143,7 @@ def test_op_roundtrip():
     )
 
 
-def test_op_roundtrip_filename(tmpdir):
+def test_op_roundtrip_filename(tmpdir) -> None:
     filename = f'{tmpdir}/op.json'
     q = cirq.LineQubit(5)
     op1 = cirq.rx(0.123).on(q)
@@ -159,7 +159,7 @@ def test_op_roundtrip_filename(tmpdir):
     assert op1 == op3
 
 
-def test_op_roundtrip_file_obj(tmpdir):
+def test_op_roundtrip_file_obj(tmpdir) -> None:
     filename = f'{tmpdir}/op.json'
     q = cirq.LineQubit(5)
     op1 = cirq.rx(0.123).on(q)
@@ -179,7 +179,7 @@ def test_op_roundtrip_file_obj(tmpdir):
         assert op1 == op3
 
 
-def test_fail_to_resolve():
+def test_fail_to_resolve() -> None:
     buffer = io.StringIO()
     buffer.write(
         """
@@ -207,7 +207,7 @@ Q0, Q1, Q2, Q3, Q4 = QUBITS
 # deprecation error in testing. It is cleaner to just turn it off than to assert
 # deprecation for each submodule.
 @mock.patch.dict(os.environ, clear='CIRQ_TESTING')
-def test_shouldnt_be_serialized_no_superfluous(mod_spec: ModuleJsonTestSpec):
+def test_shouldnt_be_serialized_no_superfluous(mod_spec: ModuleJsonTestSpec) -> None:
     # everything in the list should be ignored for a reason
     names = set(mod_spec.get_all_names())
     missing_names = set(mod_spec.should_not_be_serialized).difference(names)
@@ -223,7 +223,7 @@ def test_shouldnt_be_serialized_no_superfluous(mod_spec: ModuleJsonTestSpec):
 # deprecation error in testing. It is cleaner to just turn it off than to assert
 # deprecation for each submodule.
 @mock.patch.dict(os.environ, clear='CIRQ_TESTING')
-def test_not_yet_serializable_no_superfluous(mod_spec: ModuleJsonTestSpec):
+def test_not_yet_serializable_no_superfluous(mod_spec: ModuleJsonTestSpec) -> None:
     # everything in the list should be ignored for a reason
     names = set(mod_spec.get_all_names())
     missing_names = set(mod_spec.not_yet_serializable).difference(names)
@@ -233,7 +233,7 @@ def test_not_yet_serializable_no_superfluous(mod_spec: ModuleJsonTestSpec):
 
 
 @pytest.mark.parametrize('mod_spec', MODULE_TEST_SPECS, ids=repr)
-def test_mutually_exclusive_not_serialize_lists(mod_spec: ModuleJsonTestSpec):
+def test_mutually_exclusive_not_serialize_lists(mod_spec: ModuleJsonTestSpec) -> None:
     common = set(mod_spec.should_not_be_serialized) & set(mod_spec.not_yet_serializable)
     assert len(common) == 0, (
         f"Defined in both {mod_spec.name} 'Not yet serializable' "
@@ -242,7 +242,7 @@ def test_mutually_exclusive_not_serialize_lists(mod_spec: ModuleJsonTestSpec):
 
 
 @pytest.mark.parametrize('mod_spec', MODULE_TEST_SPECS, ids=repr)
-def test_resolver_cache_vs_should_not_serialize(mod_spec: ModuleJsonTestSpec):
+def test_resolver_cache_vs_should_not_serialize(mod_spec: ModuleJsonTestSpec) -> None:
     resolver_cache_types = set([n for (n, _) in mod_spec.get_resolver_cache_types()])
     common = set(mod_spec.should_not_be_serialized) & resolver_cache_types
 
@@ -254,7 +254,7 @@ def test_resolver_cache_vs_should_not_serialize(mod_spec: ModuleJsonTestSpec):
 
 
 @pytest.mark.parametrize('mod_spec', MODULE_TEST_SPECS, ids=repr)
-def test_resolver_cache_vs_not_yet_serializable(mod_spec: ModuleJsonTestSpec):
+def test_resolver_cache_vs_not_yet_serializable(mod_spec: ModuleJsonTestSpec) -> None:
     resolver_cache_types = set([n for (n, _) in mod_spec.get_resolver_cache_types()])
     common = set(mod_spec.not_yet_serializable) & resolver_cache_types
 
@@ -267,14 +267,14 @@ def test_resolver_cache_vs_not_yet_serializable(mod_spec: ModuleJsonTestSpec):
     )
 
 
-def test_builtins():
+def test_builtins() -> None:
     assert_json_roundtrip_works(True)
     assert_json_roundtrip_works(1)
     assert_json_roundtrip_works(1 + 2j)
     assert_json_roundtrip_works({'test': [123, 5.5], 'key2': 'asdf', '3': None, '0.0': []})
 
 
-def test_numpy():
+def test_numpy() -> None:
     x = np.ones(1)[0]
 
     assert_json_roundtrip_works(np.bool_(True))
@@ -295,7 +295,7 @@ def test_numpy():
     assert_json_roundtrip_works(np.arange(3))
 
 
-def test_pandas():
+def test_pandas() -> None:
     assert_json_roundtrip_works(
         pd.DataFrame(data=[[1, 2, 3], [4, 5, 6]], columns=['x', 'y', 'z'], index=[2, 5])
     )
@@ -320,7 +320,7 @@ def test_pandas():
     )
 
 
-def test_sympy():
+def test_sympy() -> None:
     # Raw values.
     assert_json_roundtrip_works(sympy.Symbol('theta'))
     assert_json_roundtrip_works(sympy.Integer(5))
@@ -388,7 +388,7 @@ class SBKImpl(cirq.SerializableByKey):
         return cls(name, data_list, tuple(data_tuple), data_dict)
 
 
-def test_serializable_by_key():
+def test_serializable_by_key() -> None:
     def custom_resolver(name):
         if name == 'SBKImpl':
             return SBKImpl
@@ -443,12 +443,12 @@ def _list_public_classes_for_tested_modules():
 
 
 @pytest.mark.parametrize('mod_spec,cirq_obj_name,cls', _list_public_classes_for_tested_modules())
-def test_json_test_data_coverage(mod_spec: ModuleJsonTestSpec, cirq_obj_name: str, cls):
+def test_json_test_data_coverage(mod_spec: ModuleJsonTestSpec, cirq_obj_name: str, cls) -> None:
     if cirq_obj_name in mod_spec.tested_elsewhere:
         pytest.skip("Tested elsewhere.")
 
     if cirq_obj_name in mod_spec.not_yet_serializable:
-        return pytest.xfail(reason="Not serializable (yet)")
+        pytest.xfail(reason="Not serializable (yet)")
 
     test_data_path = mod_spec.test_data_path
     rel_path = test_data_path.relative_to(REPO_ROOT)
@@ -529,12 +529,12 @@ class SerializableTypeObject:
 
 
 @pytest.mark.parametrize('mod_spec,cirq_obj_name,cls', _list_public_classes_for_tested_modules())
-def test_type_serialization(mod_spec: ModuleJsonTestSpec, cirq_obj_name: str, cls):
+def test_type_serialization(mod_spec: ModuleJsonTestSpec, cirq_obj_name: str, cls) -> None:
     if cirq_obj_name in mod_spec.tested_elsewhere:
         pytest.skip("Tested elsewhere.")
 
     if cirq_obj_name in mod_spec.not_yet_serializable:
-        return pytest.xfail(reason="Not serializable (yet)")
+        pytest.xfail(reason="Not serializable (yet)")
 
     if cls is None:
         pytest.skip(f'No serialization for None-mapped type: {cirq_obj_name}')  # pragma: no cover
@@ -556,7 +556,7 @@ def test_type_serialization(mod_spec: ModuleJsonTestSpec, cirq_obj_name: str, cl
     assert_json_roundtrip_works(sto, resolvers=test_resolvers)
 
 
-def test_invalid_type_deserialize():
+def test_invalid_type_deserialize() -> None:
     def custom_resolver(name):
         if name == 'SerializableTypeObject':
             return SerializableTypeObject
@@ -571,7 +571,7 @@ def test_invalid_type_deserialize():
         _ = cirq.read_json(json_text=factory_json, resolvers=test_resolvers)
 
 
-def test_to_from_strings():
+def test_to_from_strings() -> None:
     x_json_text = """{
   "cirq_type": "_PauliX",
   "exponent": 1.0,
@@ -584,7 +584,7 @@ def test_to_from_strings():
         cirq.read_json(io.StringIO(), json_text=x_json_text)
 
 
-def test_to_from_json_gzip():
+def test_to_from_json_gzip() -> None:
     a, b = cirq.LineQubit.range(2)
     test_circuit = cirq.Circuit(cirq.H(a), cirq.CX(a, b))
     gzip_data = cirq.to_json_gzip(test_circuit)
@@ -631,7 +631,7 @@ def assert_repr_and_json_test_data_agree(
     json_path: pathlib.Path,
     inward_only: bool,
     deprecation_deadline: str | None,
-):
+) -> None:
     if not repr_path.exists() and not json_path.exists():
         return
 
@@ -701,7 +701,7 @@ def assert_repr_and_json_test_data_agree(
     'mod_spec, abs_path',
     [(m, abs_path) for m in MODULE_TEST_SPECS for abs_path in m.all_test_data_keys()],
 )
-def test_json_and_repr_data(mod_spec: ModuleJsonTestSpec, abs_path: str):
+def test_json_and_repr_data(mod_spec: ModuleJsonTestSpec, abs_path: str) -> None:
     assert_repr_and_json_test_data_agree(
         mod_spec=mod_spec,
         repr_path=pathlib.Path(f'{abs_path}.repr'),
@@ -718,7 +718,7 @@ def test_json_and_repr_data(mod_spec: ModuleJsonTestSpec, abs_path: str):
     )
 
 
-def test_pathlib_paths(tmpdir):
+def test_pathlib_paths(tmpdir) -> None:
     path = pathlib.Path(tmpdir) / 'op.json'
     cirq.to_json(cirq.X, path)
     assert cirq.read_json(path) == cirq.X
@@ -746,7 +746,7 @@ def test_dataclass_json_dict() -> None:
     assert_json_roundtrip_works(my_dc, resolvers=[custom_resolver, *cirq.DEFAULT_RESOLVERS])
 
 
-def test_numpy_values():
+def test_numpy_values() -> None:
     assert (
         cirq.to_json({'value': np.array(1)})
         == """{
@@ -755,7 +755,7 @@ def test_numpy_values():
     )
 
 
-def test_basic_time_assertions():
+def test_basic_time_assertions() -> None:
     naive_dt = datetime.datetime.now()
     utc_dt = naive_dt.astimezone(datetime.timezone.utc)
     assert naive_dt.timestamp() == utc_dt.timestamp()
@@ -768,7 +768,7 @@ def test_basic_time_assertions():
     assert naive_dt == re_naive, 'works, as long as you called fromtimestamp from the same timezone'
 
 
-def test_datetime():
+def test_datetime() -> None:
     naive_dt = datetime.datetime.now()
 
     re_naive_dt = cirq.read_json(json_text=cirq.to_json(naive_dt))
@@ -793,7 +793,7 @@ class _TestAttrsClas:
     x: int
 
 
-def test_attrs_json_dict():
+def test_attrs_json_dict() -> None:
     obj = _TestAttrsClas('test', x=123)
     js = json_serialization.attrs_json_dict(obj)
     assert js == {'name': 'test', 'x': 123}

cirq/protocols/kraus_protocol.py

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Iterable
 from types import NotImplementedType
-from typing import Any, Protocol, Sequence, TypeVar
+from typing import Any, Protocol, TypeVar
 
 import numpy as np
 
@@ -31,7 +32,7 @@ from cirq.protocols.unitary_protocol import unitary
 
 # This is a special indicator value used by the channel method to determine
 # whether or not the caller provided a 'default' argument. It must be of type
-# Sequence[np.ndarray] to ensure the method has the correct type signature in
+# Iterable[np.ndarray] to ensure the method has the correct type signature in
 # that case. It is checked for using `is`, so it won't have a false positive
 # if the user provides a different (np.array([]),) value.
 RaiseTypeErrorIfNotProvided: tuple[np.ndarray] = (np.array([]),)
@@ -44,7 +45,7 @@ class SupportsKraus(Protocol):
     """An object that may be describable as a quantum channel."""
 
     @doc_private
-    def _kraus_(self) -> Sequence[np.ndarray] | NotImplementedType:
+    def _kraus_(self) -> Iterable[np.ndarray] | NotImplementedType:
         r"""A list of Kraus matrices describing the quantum channel.
 
         These matrices are the terms in the operator sum representation of a

cirq/protocols/kraus_protocol_test.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable, Sequence
+from collections.abc import Iterable
 
 import numpy as np
 import pytest
@@ -42,7 +42,7 @@ def test_kraus_no_methods() -> None:
     assert not cirq.has_kraus(NoMethod())
 
 
-def assert_not_implemented(val):
+def assert_not_implemented(val) -> None:
     with pytest.raises(TypeError, match='returned NotImplemented'):
         _ = cirq.kraus(val)
 
@@ -89,7 +89,7 @@ def test_explicit_kraus() -> None:
     c = (a0, a1)
 
     class ReturnsKraus:
-        def _kraus_(self) -> Sequence[np.ndarray]:
+        def _kraus_(self) -> Iterable[np.ndarray]:
             return c
 
     assert cirq.kraus(ReturnsKraus()) is c
@@ -174,7 +174,7 @@ def test_has_kraus_when_decomposed(decomposed_cls) -> None:
     assert not cirq.has_kraus(op, allow_decompose=False)
 
 
-def test_strat_kraus_from_apply_channel_returns_none():
+def test_strat_kraus_from_apply_channel_returns_none() -> None:
     # Remove _kraus_ and _apply_channel_ methods
     class NoApplyChannelReset(cirq.ResetChannel):
         def _kraus_(self):
@@ -228,7 +228,7 @@ def test_kraus_fallback_to_apply_channel(channel_cls, params) -> None:
     np.testing.assert_allclose(actual_super, expected_super, atol=1e-8)
 
 
-def test_reset_channel_kraus_apply_channel_consistency():
+def test_reset_channel_kraus_apply_channel_consistency() -> None:
     Reset = cirq.ResetChannel
     # Original gate
     gate = Reset()
@@ -236,7 +236,7 @@ def test_reset_channel_kraus_apply_channel_consistency():
     cirq.testing.assert_consistent_channel(gate)
 
     # Remove _kraus_ method
-    class NoKrausReset(Reset):
+    class NoKrausReset(cirq.ResetChannel):
         def _kraus_(self):
             return NotImplemented
 
@@ -245,7 +245,7 @@ def test_reset_channel_kraus_apply_channel_consistency():
     np.testing.assert_allclose(cirq.kraus(gate), cirq.kraus(gate_no_kraus), atol=1e-8)
 
 
-def test_kraus_channel_with_has_unitary():
+def test_kraus_channel_with_has_unitary() -> None:
     """CZSWAP has no unitary dunder method but has_unitary returns True."""
     op = cirq.CZSWAP.on(cirq.q(1), cirq.q(2))
     channels = cirq.kraus(op)

cirq/protocols/measurement_key_protocol.py

@@ -16,8 +16,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from types import NotImplementedType
-from typing import Any, Mapping, Protocol, TYPE_CHECKING
+from typing import Any, overload, Protocol, TYPE_CHECKING, TypeVar
 
 from cirq import value
 from cirq._doc import doc_private
@@ -25,6 +26,9 @@ from cirq._doc import doc_private
 if TYPE_CHECKING:
     import cirq
 
+T = TypeVar('T')
+TDefault = TypeVar('TDefault')
+
 # This is a special indicator value used by the inverse method to determine
 # whether or not the caller provided a 'default' argument.
 RaiseTypeErrorIfNotProvided: Any = ([],)
@@ -104,7 +108,17 @@ class SupportsMeasurementKey(Protocol):
         """
 
 
-def measurement_key_obj(val: Any, default: Any = RaiseTypeErrorIfNotProvided):
+@overload
+def measurement_key_obj(val: Any) -> cirq.MeasurementKey:
+    pass
+
+
+@overload
+def measurement_key_obj(val: Any, default: TDefault) -> cirq.MeasurementKey | TDefault:
+    pass
+
+
+def measurement_key_obj(val, default=RaiseTypeErrorIfNotProvided):
     """Get the single measurement key object for the given value.
 
     Args:
@@ -138,7 +152,17 @@ def measurement_key_obj(val: Any, default: Any = RaiseTypeErrorIfNotProvided):
     raise TypeError(f"Object of type '{type(val)}' had no measurement keys.")
 
 
-def measurement_key_name(val: Any, default: Any = RaiseTypeErrorIfNotProvided):
+@overload
+def measurement_key_name(val: Any) -> str:
+    pass
+
+
+@overload
+def measurement_key_name(val: Any, default: TDefault) -> str | TDefault:
+    pass
+
+
+def measurement_key_name(val, default=RaiseTypeErrorIfNotProvided):
     """Get the single measurement key for the given value.
 
     Args:
@@ -278,7 +302,7 @@ def is_measurement(val: Any) -> bool:
     return keys is not NotImplemented and bool(keys)
 
 
-def with_measurement_key_mapping(val: Any, key_map: Mapping[str, str]):
+def with_measurement_key_mapping(val: T, key_map: Mapping[str, str]) -> T:
     """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
@@ -288,7 +312,7 @@ def with_measurement_key_mapping(val: Any, key_map: Mapping[str, str]):
     return NotImplemented if getter is None else getter(key_map)
 
 
-def with_key_path(val: Any, path: tuple[str, ...]):
+def with_key_path(val: T, path: tuple[str, ...]) -> T:
     """Adds the path to the target's measurement keys.
 
     The path usually refers to an identifier or a list of identifiers from a subcircuit that
@@ -299,7 +323,7 @@ def with_key_path(val: Any, path: tuple[str, ...]):
     return NotImplemented if getter is None else getter(path)
 
 
-def with_key_path_prefix(val: Any, prefix: tuple[str, ...]):
+def with_key_path_prefix(val: T, prefix: tuple[str, ...]) -> T:
     """Prefixes the path to the target's measurement keys.
 
     The path usually refers to an identifier or a list of identifiers from a subcircuit that
@@ -315,8 +339,8 @@ def with_key_path_prefix(val: Any, prefix: tuple[str, ...]):
 
 
 def with_rescoped_keys(
-    val: Any, path: tuple[str, ...], bindable_keys: frozenset[cirq.MeasurementKey] | None = None
-):
+    val: T, path: tuple[str, ...], bindable_keys: frozenset[cirq.MeasurementKey] | None = None
+) -> T:
     """Rescopes any measurement and control keys to the provided path, given the existing keys.
 
     The path usually refers to an identifier or a list of identifiers from a subcircuit that

cirq/protocols/mixture_protocol.py

@@ -16,8 +16,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from types import NotImplementedType
-from typing import Any, Protocol, Sequence
+from typing import Any, Protocol
 
 import numpy as np
 
@@ -148,7 +149,7 @@ def has_mixture(val: Any, *, allow_decompose: bool = True) -> bool:
     return mixture(val, None) is not None
 
 
-def validate_mixture(supports_mixture: SupportsMixture):
+def validate_mixture(supports_mixture: SupportsMixture) -> None:
     """Validates that the mixture's tuple are valid probabilities."""
     mixture_tuple = mixture(supports_mixture, None)
     if mixture_tuple is None:

cirq/protocols/mixture_protocol_test.py

@@ -93,7 +93,7 @@ class ReturnsMixtureOfReturnsUnitary:
         (ReturnsMixtureOfReturnsUnitary(), ((0.4, np.eye(2)), (0.6, np.eye(2)))),
     ),
 )
-def test_objects_with_mixture(val, mixture):
+def test_objects_with_mixture(val, mixture) -> None:
     expected_keys, expected_values = zip(*mixture)
     keys, values = zip(*cirq.mixture(val))
     np.testing.assert_almost_equal(keys, expected_keys)
@@ -107,7 +107,7 @@ def test_objects_with_mixture(val, mixture):
 @pytest.mark.parametrize(
     'val', (NoMethod(), ReturnsNotImplemented(), ReturnsNotImplementedUnitary())
 )
-def test_objects_with_no_mixture(val):
+def test_objects_with_no_mixture(val) -> None:
     with pytest.raises(TypeError, match="mixture"):
         _ = cirq.mixture(val)
     assert cirq.mixture(val, None) is None
@@ -116,7 +116,7 @@ def test_objects_with_no_mixture(val):
     assert cirq.mixture(val, default) == default
 
 
-def test_has_mixture():
+def test_has_mixture() -> None:
     assert cirq.has_mixture(ReturnsValidTuple())
     assert not cirq.has_mixture(ReturnsNotImplemented())
     assert cirq.has_mixture(ReturnsMixtureButNoHasMixture())
@@ -124,7 +124,7 @@ def test_has_mixture():
     assert not cirq.has_mixture(ReturnsNotImplementedUnitary())
 
 
-def test_valid_mixture():
+def test_valid_mixture() -> None:
     cirq.validate_mixture(ReturnsValidTuple())
 
 
@@ -136,11 +136,11 @@ def test_valid_mixture():
         (ReturnsGreaterThanUnityProbability(), 'greater than 1'),
     ),
 )
-def test_invalid_mixture(val, message):
+def test_invalid_mixture(val, message) -> None:
     with pytest.raises(ValueError, match=message):
         cirq.validate_mixture(val)
 
 
-def test_missing_mixture():
+def test_missing_mixture() -> None:
     with pytest.raises(TypeError, match='_mixture_'):
-        cirq.validate_mixture(NoMethod)
+        cirq.validate_mixture(NoMethod)  # type: ignore[arg-type]

cirq/protocols/mul_protocol_test.py

@@ -52,7 +52,7 @@ class MulSevenRMulEight:
         return 8
 
 
-def test_equivalent_to_builtin_mul():
+def test_equivalent_to_builtin_mul() -> None:
     test_vals = [
         0,
         1,
@@ -76,14 +76,14 @@ def test_equivalent_to_builtin_mul():
             c = cirq.mul(a, b, default=None)
             if c is None:
                 with pytest.raises(TypeError):
-                    _ = a * b
+                    _ = a * b  # type: ignore[operator]
                 with pytest.raises(TypeError):
                     _ = cirq.mul(a, b)
             else:
-                assert c == a * b
+                assert c == a * b  # type: ignore[operator]
 
 
-def test_symbol_special_case():
+def test_symbol_special_case() -> None:
     x = sympy.Symbol('x')
     assert cirq.mul(x, 1.0) is x
     assert cirq.mul(1.0, x) is x

cirq/protocols/phase_protocol.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Protocol, TypeVar
+from typing import Any, overload, Protocol, TypeVar
 
 # This is a special value to indicate that a type error should be returned.
 # This is used within phase_by to raise an error if no underlying
@@ -24,6 +24,7 @@ from cirq._doc import doc_private
 RaiseTypeErrorIfNotProvided: Any = ([],)
 
 TDefault = TypeVar('TDefault')
+T = TypeVar('T')
 
 
 class SupportsPhase(Protocol):
@@ -47,9 +48,17 @@ class SupportsPhase(Protocol):
         """
 
 
-def phase_by(
-    val: Any, phase_turns: float, qubit_index: int, default: TDefault = RaiseTypeErrorIfNotProvided
-):
+@overload
+def phase_by(val: T, phase_turns: float, qubit_index: int) -> T:
+    pass
+
+
+@overload
+def phase_by(val: T, phase_turns: float, qubit_index: int, default: TDefault) -> T | TDefault:
+    pass
+
+
+def phase_by(val, phase_turns, qubit_index, default=RaiseTypeErrorIfNotProvided):
     """Returns a phased version of the effect.
 
     For example, an X gate phased by 90 degrees would be a Y gate.

cirq/protocols/pow_protocol.py

@@ -16,7 +16,8 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, overload, TYPE_CHECKING, TypeVar
+from collections.abc import Callable
+from typing import Any, overload, TYPE_CHECKING, TypeVar
 
 if TYPE_CHECKING:
     import cirq

cirq/protocols/pow_protocol_test.py

@@ -34,7 +34,7 @@ class ReturnsExponent:
 
 
 @pytest.mark.parametrize('val', (NoMethod(), 'text', object(), ReturnsNotImplemented()))
-def test_powerless(val):
+def test_powerless(val) -> None:
     assert cirq.pow(val, 5, None) is None
     assert cirq.pow(val, 2, NotImplemented) is NotImplemented
 
@@ -42,13 +42,13 @@ def test_powerless(val):
     assert cirq.pow(val, 1, None) is None
 
 
-def test_pow_error():
+def test_pow_error() -> None:
     with pytest.raises(TypeError, match="returned NotImplemented"):
-        _ = cirq.pow(ReturnsNotImplemented(), 3)
+        _ = cirq.pow(ReturnsNotImplemented(), 3)  # type: ignore[call-overload]
     with pytest.raises(TypeError, match="no __pow__ method"):
-        _ = cirq.pow(NoMethod(), 3)
+        _ = cirq.pow(NoMethod(), 3)  # type: ignore[call-overload]
 
 
 @pytest.mark.parametrize('val,exponent,out', ((ReturnsExponent(), 2, 2), (1, 2, 1), (2, 3, 8)))
-def test_pow_with_result(val, exponent, out):
+def test_pow_with_result(val, exponent, out) -> None:
     assert cirq.pow(val, exponent) == cirq.pow(val, exponent, default=None) == val**exponent == out

cirq/protocols/qasm.py

@@ -15,8 +15,9 @@
 from __future__ import annotations
 
 import string
+from collections.abc import Iterable, Mapping
 from types import NotImplementedType
-from typing import Any, Iterable, Mapping, Protocol, TYPE_CHECKING, TypeVar
+from typing import Any, Protocol, TYPE_CHECKING, TypeVar
 
 from cirq import ops
 from cirq._doc import doc_private

cirq/protocols/qid_shape_protocol.py

@@ -14,8 +14,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from types import NotImplementedType
-from typing import Any, Protocol, Sequence, TypeVar
+from typing import Any, Protocol, TypeVar
 
 from cirq import ops
 from cirq._doc import doc_private, document