Add return-type to public functions, mostly tests part 5 (#7691)

No change in the effective code. A batch of ~50 files.

Modified files pass ruff check --select=ANN201

Notable changes:

- Add no-op method `Operation.__pow__` to support static type check
- Replace `list[Qid]` --> `Sequence[Qid]` which is covariant
- Make `KakDecomposition` and `FSimGate` match the `SupportsUnitary`
type
- Loosen return type of the `_kraus_` method from `Sequence` to
`Iterable`

Partially implements #4393

Author: pavoljuhas

cirq-core/cirq/_doc.py

@@ -15,12 +15,13 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TypeVar
 
+T = TypeVar('T')
 RECORDED_CONST_DOCS: dict[int, str] = {}
 
 
-def document(value: Any, doc_string: str = ''):
+def document(value: T, doc_string: str = '') -> T:
     """Stores documentation details about the given value.
 
     This method is used to associate a docstring with global constants. It is
@@ -64,7 +65,7 @@ def document(value: Any, doc_string: str = ''):
 _DOC_PRIVATE = "_tf_docs_doc_private"
 
 
-def doc_private(obj):
+def doc_private(obj: T) -> T:
     """A decorator: Generates docs for private methods/functions.
 
     For example:

cirq-core/cirq/linalg/decompositions.py

@@ -472,6 +472,9 @@ def __repr__(self) -> str:
             f'    global_phase={self.global_phase!r})'
         )
 
+    def _has_unitary_(self) -> bool:
+        return True
+
     def _unitary_(self) -> np.ndarray:
         """Returns the decomposition's two-qubit unitary matrix.
 

cirq-core/cirq/linalg/decompositions_test.py

@@ -259,6 +259,7 @@ def test_kak_plot_empty() -> None:
 )
 def test_kak_decomposition(target) -> None:
     kak = cirq.kak_decomposition(target)
+    assert cirq.has_unitary(kak)
     np.testing.assert_allclose(cirq.unitary(kak), target, atol=1e-8)
 
 

cirq-core/cirq/ops/fourier_transform_test.py

@@ -107,9 +107,7 @@ def test_qft() -> None:
 
     arr = np.array([[1, 1, 1, 1], [1, -1j, -1, 1j], [1, -1, 1, -1], [1, 1j, -1, -1j]]) / 2
     np.testing.assert_allclose(
-        cirq.unitary(cirq.qft(*cirq.LineQubit.range(2)) ** -1),  # type: ignore[operator]
-        arr,  # type: ignore[arg-type]
-        atol=1e-8,
+        cirq.unitary(cirq.qft(*cirq.LineQubit.range(2)) ** -1), arr, atol=1e-8
     )
 
     for k in range(4):
@@ -121,7 +119,7 @@ def test_qft() -> None:
 
 def test_inverse() -> None:
     a, b, c = cirq.LineQubit.range(3)
-    assert cirq.qft(a, b, c, inverse=True) == cirq.qft(a, b, c) ** -1  # type: ignore[operator]
+    assert cirq.qft(a, b, c, inverse=True) == cirq.qft(a, b, c) ** -1
     assert cirq.qft(a, b, c, inverse=True, without_reverse=True) == cirq.inverse(
         cirq.qft(a, b, c, without_reverse=True)
     )

cirq-core/cirq/ops/raw_types.py

@@ -528,6 +528,9 @@ def _num_qubits_(self) -> int:
     def _qid_shape_(self) -> tuple[int, ...]:
         return protocols.qid_shape(self.qubits)
 
+    def __pow__(self, exponent: Any) -> Operation:
+        return NotImplemented  # pragma: no cover
+
     @abc.abstractmethod
     def with_qubits(self, *new_qubits: cirq.Qid) -> cirq.Operation:
         """Returns the same operation, but applied to different qubits.

cirq-core/cirq/protocols/apply_unitary_protocol_test.py

@@ -240,6 +240,7 @@ def assert_is_swap_simple(val: cirq.SupportsConsistentApplyUnitary) -> None:
             op_indices, tuple(qid_shape[i] for i in op_indices)
         )
         sub_result = val._apply_unitary_(sub_args)
+        assert isinstance(sub_result, np.ndarray)
         result = _incorporate_result_into_target(args, sub_args, sub_result)
         np.testing.assert_allclose(result, expected, atol=1e-8)
 
@@ -258,6 +259,7 @@ def assert_is_swap(val: cirq.SupportsConsistentApplyUnitary) -> None:
             op_indices, tuple(qid_shape[i] for i in op_indices)
         )
         sub_result = val._apply_unitary_(sub_args)
+        assert isinstance(sub_result, np.ndarray)
         result = _incorporate_result_into_target(args, sub_args, sub_result)
         np.testing.assert_allclose(result, expected, atol=1e-8, verbose=True)
 

cirq-core/cirq/protocols/kraus_protocol.py

@@ -18,7 +18,7 @@
 
 import warnings
 from types import NotImplementedType
-from typing import Any, Protocol, Sequence, TypeVar
+from typing import Any, Iterable, Protocol, TypeVar
 
 import numpy as np
 
@@ -31,7 +31,7 @@
 
 # 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 +44,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-core/cirq/protocols/kraus_protocol_test.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable, Sequence
+from typing import Iterable
 
 import numpy as np
 import pytest
@@ -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

cirq-core/cirq/protocols/unitary_protocol.py

@@ -38,7 +38,7 @@ class SupportsUnitary(Protocol):
     """An object that may be describable by a unitary matrix."""
 
     @doc_private
-    def _unitary_(self) -> np.ndarray | NotImplementedType:
+    def _unitary_(self) -> np.ndarray | NotImplementedType | None:
         """A unitary matrix describing this value, e.g. the matrix of a gate.
 
         This method is used by the global `cirq.unitary` method. If this method

cirq-core/cirq/testing/consistent_resolve_parameters.py

@@ -21,7 +21,7 @@
 import cirq
 
 
-def assert_consistent_resolve_parameters(val: Any):
+def assert_consistent_resolve_parameters(val: Any) -> None:
     names = cirq.parameter_names(val)
     symbols = cirq.parameter_symbols(val)