cirq-core 1.6.0.dev20250702182429__py3-none-any.whl → 1.6.0.dev20250704055727__py3-none-any.whl

cirq/contrib/acquaintance/gates_test.py

@@ -27,7 +27,7 @@ import cirq.contrib.acquaintance as cca
 import cirq.testing as ct
 
 
-def test_acquaintance_gate_repr():
+def test_acquaintance_gate_repr() -> None:
     assert (
         repr(cca.AcquaintanceOpportunityGate(2))
         == 'cirq.contrib.acquaintance.AcquaintanceOpportunityGate'
@@ -35,7 +35,7 @@ def test_acquaintance_gate_repr():
     )
 
 
-def test_acquaintance_gate_text_diagram_info():
+def test_acquaintance_gate_text_diagram_info() -> None:
     qubits = [cirq.NamedQubit(s) for s in 'xyz']
     circuit = cirq.Circuit([cirq.Moment([cca.acquaint(*qubits)])])
     actual_text_diagram = circuit.to_text_diagram().strip()
@@ -49,11 +49,11 @@ z: ───█───
     assert actual_text_diagram == expected_text_diagram
 
 
-def test_acquaintance_gate_unknown_qubit_count():
+def test_acquaintance_gate_unknown_qubit_count() -> None:
     assert cirq.circuit_diagram_info(cca.acquaint, default=None) is None
 
 
-def test_swap_network_gate():
+def test_swap_network_gate() -> None:
     qubits = tuple(cirq.NamedQubit(s) for s in alphabet)
 
     acquaintance_size = 3
@@ -125,7 +125,7 @@ f: ───╱1╲─────────╱1╲─────────
 @pytest.mark.parametrize(
     'part_lens', [tuple(randint(1, 3) for _ in range(randint(2, 10))) for _ in range(3)]
 )
-def test_acquaint_part_pairs(part_lens):
+def test_acquaint_part_pairs(part_lens) -> None:
     parts = []
     n_qubits = 0
     for part_len in part_lens:
@@ -154,13 +154,13 @@ acquaintance_sizes += tuple(range(5))
         )
     ),
 )
-def test_swap_network_gate_permutation(part_lens, acquaintance_size):
+def test_swap_network_gate_permutation(part_lens, acquaintance_size) -> None:
     n_qubits = sum(part_lens)
     swap_network_gate = cca.SwapNetworkGate(part_lens, acquaintance_size)
     cca.testing.assert_permutation_decomposition_equivalence(swap_network_gate, n_qubits)
 
 
-def test_swap_network_gate_from_ops():
+def test_swap_network_gate_from_ops() -> None:
     n_qubits = 10
     qubits = cirq.LineQubit.range(n_qubits)
     part_lens = (1, 2, 1, 3, 3)
@@ -198,7 +198,7 @@ def test_swap_network_gate_from_ops():
     cirq.testing.assert_has_diagram(circuit, expected_diagram)
 
 
-def test_swap_network_decomposition():
+def test_swap_network_decomposition() -> None:
     qubits = cirq.LineQubit.range(8)
     swap_network_gate = cca.SwapNetworkGate((4, 4), 5)
     operations = cirq.decompose_once_with_qubits(swap_network_gate, qubits)
@@ -223,7 +223,7 @@ def test_swap_network_decomposition():
     ct.assert_has_diagram(circuit, expected_text_diagram)
 
 
-def test_swap_network_init_error():
+def test_swap_network_init_error() -> None:
     with pytest.raises(ValueError):
         cca.SwapNetworkGate(())
     with pytest.raises(ValueError):
@@ -237,7 +237,7 @@ part_lens_and_acquaintance_sizes = [
 
 
 @pytest.mark.parametrize('part_lens, acquaintance_size', part_lens_and_acquaintance_sizes)
-def test_swap_network_permutation(part_lens, acquaintance_size):
+def test_swap_network_permutation(part_lens, acquaintance_size) -> None:
     n_qubits = sum(part_lens)
     gate = cca.SwapNetworkGate(part_lens, acquaintance_size)
 
@@ -246,7 +246,7 @@ def test_swap_network_permutation(part_lens, acquaintance_size):
 
 
 @pytest.mark.parametrize('part_lens, acquaintance_size', part_lens_and_acquaintance_sizes)
-def test_swap_network_repr(part_lens, acquaintance_size):
+def test_swap_network_repr(part_lens, acquaintance_size) -> None:
     gate = cca.SwapNetworkGate(part_lens, acquaintance_size)
     ct.assert_equivalent_repr(gate)
 
@@ -260,16 +260,18 @@ class OtherOperation(cirq.Operation):
         return self._qubits
 
     def with_qubits(self, *new_qubits: cirq.Qid) -> OtherOperation:
-        return type(self)(self._qubits)
+        return type(self)(new_qubits)
 
     def __eq__(self, other):
         return isinstance(other, type(self)) and self.qubits == other.qubits
 
 
-def test_get_acquaintance_size():
+def test_get_acquaintance_size() -> None:
+    qubits: Sequence[cirq.Qid]
     qubits = cirq.LineQubit.range(8)
+    op: cirq.Operation
     op = OtherOperation(qubits)
-    assert op.with_qubits(qubits) == op
+    assert op.with_qubits(*qubits) == op
     assert cca.get_acquaintance_size(op) == 0
 
     for s, _ in enumerate(qubits):
@@ -301,7 +303,7 @@ def test_get_acquaintance_size():
     assert cca.get_acquaintance_size(op) == 0
 
 
-def test_operations_to_part_lens():
+def test_operations_to_part_lens() -> None:
     qubits = cirq.LineQubit.range(6)
     ops = [cirq.CZ(*qubits[1:3]), cirq.XX(*qubits[3:5])]
     part_lens = cca.gates.operations_to_part_lens(qubits, ops)
@@ -319,7 +321,7 @@ def test_operations_to_part_lens():
 @pytest.mark.parametrize(
     'part_len_sets', [set(tuple(randint(1, 5) for _ in range(randint(2, 7))) for _ in range(5))]
 )
-def test_swap_network_gate_equality(part_len_sets):
+def test_swap_network_gate_equality(part_len_sets) -> None:
     acquaintance_sizes = [None, 0, 1, 2, 3]
     swap_gates = [cirq.SWAP, cirq.CNOT]
     equals_tester = ct.EqualsTester()

cirq/contrib/acquaintance/inspection_utils.py

@@ -51,7 +51,9 @@ class LogicalAnnotator(ExecutionStrategy):
         yield AcquaintanceOperation(qubits, indices)
 
 
-def get_acquaintance_dag(strategy: cirq.Circuit, initial_mapping: LogicalMapping):
+def get_acquaintance_dag(
+    strategy: cirq.Circuit, initial_mapping: LogicalMapping
+) -> circuitdag.CircuitDag:
     strategy = strategy.copy()
     expose_acquaintance_gates(strategy)
     LogicalAnnotator(initial_mapping)(strategy)
@@ -66,9 +68,11 @@ def get_acquaintance_dag(strategy: cirq.Circuit, initial_mapping: LogicalMapping
 
 def get_logical_acquaintance_opportunities(
     strategy: cirq.Circuit, initial_mapping: LogicalMapping
-) -> set[frozenset[LogicalIndex]]:
+) -> set[frozenset[int]] | set[frozenset[cirq.Qid]]:
     acquaintance_dag = get_acquaintance_dag(strategy, initial_mapping)
     logical_acquaintance_opportunities = set()
     for op in acquaintance_dag.all_operations():
-        logical_acquaintance_opportunities.add(frozenset(op.logical_indices))
+        logical_acquaintance_opportunities.add(
+            frozenset(op.logical_indices)  # type: ignore[attr-defined]
+        )
     return logical_acquaintance_opportunities

cirq/contrib/acquaintance/permutation_test.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import random
+from typing import Iterator
 
 import pytest
 
@@ -23,7 +24,7 @@ import cirq.contrib.acquaintance as cca
 import cirq.testing as ct
 
 
-def test_swap_permutation_gate():
+def test_swap_permutation_gate() -> None:
     no_decomp = lambda op: (isinstance(op, cirq.GateOperation) and op.gate == cirq.SWAP)
     a, b = cirq.NamedQubit('a'), cirq.NamedQubit('b')
     gate = cca.SwapPermutationGate()
@@ -42,7 +43,7 @@ def test_swap_permutation_gate():
         cirq.commutes(gate, cirq.CCZ)
 
 
-def test_validate_permutation_errors():
+def test_validate_permutation_errors() -> None:
     validate_permutation = cca.PermutationGate.validate_permutation
     validate_permutation({})
 
@@ -59,7 +60,7 @@ def test_validate_permutation_errors():
     assert cirq.circuit_diagram_info(gate, default=None) is None
 
 
-def test_diagram():
+def test_diagram() -> None:
     gate = cca.SwapPermutationGate()
     a, b = cirq.NamedQubit('a'), cirq.NamedQubit('b')
     circuit = cirq.Circuit([gate(a, b)])
@@ -72,17 +73,19 @@ b: ───1↦0───
     assert actual_text_diagram == expected_text_diagram
 
 
-def test_update_mapping():
+def test_update_mapping() -> None:
     gate = cca.SwapPermutationGate()
     a, b, c = (cirq.NamedQubit(s) for s in 'abc')
+    mapping: dict[cirq.Qid, int]
     mapping = {s: i for i, s in enumerate((a, b, c))}
     ops = [gate(a, b), gate(b, c)]
     cca.update_mapping(mapping, ops)
     assert mapping == {a: 1, b: 2, c: 0}
 
 
-def test_get_logical_operations():
+def test_get_logical_operations() -> None:
     a, b, c, d = qubits = cirq.LineQubit.range(4)
+    mapping: dict[cirq.Qid, cirq.Qid]
     mapping = dict(zip(qubits, qubits))
     operations = [
         cirq.ZZ(a, b),
@@ -106,7 +109,7 @@ def test_get_logical_operations():
         for n_elements in (random.randint(5, 20) for _ in range(20))
     ),
 )
-def test_linear_permutation_gate(n_elements, n_permuted):
+def test_linear_permutation_gate(n_elements, n_permuted) -> None:
     qubits = cirq.LineQubit.range(n_elements)
     elements = tuple(range(n_elements))
     elements_to_permute = random.sample(elements, n_permuted)
@@ -116,16 +119,18 @@ def test_linear_permutation_gate(n_elements, n_permuted):
     gate = cca.LinearPermutationGate(n_elements, permutation)
     ct.assert_equivalent_repr(gate)
     assert gate.permutation() == permutation
+    mapping: dict[cirq.Qid, int]
     mapping = dict(zip(qubits, elements))
     for swap in cirq.flatten_op_tree(cirq.decompose_once_with_qubits(gate, qubits)):
         assert isinstance(swap, cirq.GateOperation)
+        assert isinstance(swap.gate, cca.PermutationGate)
         swap.gate.update_mapping(mapping, swap.qubits)
     for i in range(n_elements):
         p = permutation.get(elements[i], i)
         assert mapping.get(qubits[p], elements[i]) == i
 
 
-def random_equal_permutations(n_perms, n_items, prob):
+def random_equal_permutations(n_perms, n_items, prob) -> list[dict[int, int]]:
     indices_to_permute = [i for i in range(n_items) if random.random() <= prob]
     permuted_indices = random.sample(indices_to_permute, len(indices_to_permute))
     base_permutation = dict(zip(indices_to_permute, permuted_indices))
@@ -138,7 +143,9 @@ def random_equal_permutations(n_perms, n_items, prob):
     return permutations
 
 
-def random_permutation_equality_groups(n_groups, n_perms_per_group, n_items, prob):
+def random_permutation_equality_groups(
+    n_groups, n_perms_per_group, n_items, prob
+) -> Iterator[list[dict[int, int]]]:
     fingerprints = set()
     for _ in range(n_groups):
         perms = random_equal_permutations(n_perms_per_group, n_items, prob)
@@ -150,7 +157,7 @@ def random_permutation_equality_groups(n_groups, n_perms_per_group, n_items, pro
 
 
 @pytest.mark.parametrize('permutation_sets', [random_permutation_equality_groups(5, 3, 10, 0.5)])
-def test_linear_permutation_gate_equality(permutation_sets):
+def test_linear_permutation_gate_equality(permutation_sets) -> None:
     swap_gates = [cirq.SWAP, cirq.CNOT]
     equals_tester = ct.EqualsTester()
     for swap_gate in swap_gates:
@@ -163,7 +170,7 @@ def test_linear_permutation_gate_equality(permutation_sets):
             )
 
 
-def test_linear_permutation_gate_pow_not_implemented():
+def test_linear_permutation_gate_pow_not_implemented() -> None:
     permutation_gate = cca.LinearPermutationGate(3, {0: 1, 1: 2, 2: 0})
 
     assert permutation_gate.__pow__(0) is NotImplemented
@@ -183,7 +190,7 @@ def test_linear_permutation_gate_pow_not_implemented():
         (4, {0: 3, 1: 2, 2: 1, 3: 0}),
     ],
 )
-def test_linear_permutation_gate_pow_identity(num_qubits, permutation):
+def test_linear_permutation_gate_pow_identity(num_qubits, permutation) -> None:
     permutation_gate = cca.LinearPermutationGate(num_qubits, permutation)
 
     assert permutation_gate**1 == permutation_gate
@@ -199,7 +206,7 @@ def test_linear_permutation_gate_pow_identity(num_qubits, permutation):
         (4, {0: 3, 1: 2, 2: 1, 3: 0}, {0: 3, 1: 2, 2: 1, 3: 0}),
     ],
 )
-def test_linear_permutation_gate_pow_inverse(num_qubits, permutation, inverse):
+def test_linear_permutation_gate_pow_inverse(num_qubits, permutation, inverse) -> None:
     permutation_gate = cca.LinearPermutationGate(num_qubits, permutation)
     inverse_gate = cca.LinearPermutationGate(num_qubits, inverse)
 
@@ -207,7 +214,7 @@ def test_linear_permutation_gate_pow_inverse(num_qubits, permutation, inverse):
     assert cirq.inverse(permutation_gate) == inverse_gate
 
 
-def test_display_mapping():
+def test_display_mapping() -> None:
     indices = [4, 2, 0, 1, 3]
     qubits = cirq.LineQubit.range(len(indices))
     circuit = cca.complete_acquaintance_strategy(qubits, 2)
@@ -238,7 +245,7 @@ def test_display_mapping():
         for _ in range(4)
     ],
 )
-def test_return_to_initial_mapping(circuit):
+def test_return_to_initial_mapping(circuit) -> None:
     qubits = sorted(circuit.all_qubits())
     cca.return_to_initial_mapping(circuit)
     initial_mapping = {q: i for i, q in enumerate(qubits)}
@@ -247,7 +254,7 @@ def test_return_to_initial_mapping(circuit):
     assert mapping == initial_mapping
 
 
-def test_uses_consistent_swap_gate():
+def test_uses_consistent_swap_gate() -> None:
     a, b = cirq.LineQubit.range(2)
     circuit = cirq.Circuit([cca.SwapPermutationGate()(a, b), cca.SwapPermutationGate()(a, b)])
     assert cca.uses_consistent_swap_gate(circuit, cirq.SWAP)
@@ -264,7 +271,7 @@ def test_uses_consistent_swap_gate():
     assert not cca.uses_consistent_swap_gate(circuit, cirq.CZ)
 
 
-def test_swap_gate_eq():
+def test_swap_gate_eq() -> None:
     assert cca.SwapPermutationGate() == cca.SwapPermutationGate(cirq.SWAP)
     assert cca.SwapPermutationGate() != cca.SwapPermutationGate(cirq.CZ)
     assert cca.SwapPermutationGate(cirq.CZ) == cca.SwapPermutationGate(cirq.CZ)
@@ -278,5 +285,5 @@ def test_swap_gate_eq():
         cca.SwapPermutationGate(cirq.CZ),
     ],
 )
-def test_swap_gate_repr(gate):
+def test_swap_gate_repr(gate) -> None:
     cirq.testing.assert_equivalent_repr(gate)

cirq/contrib/acquaintance/shift_swap_network_test.py

@@ -23,7 +23,7 @@ import cirq
 import cirq.contrib.acquaintance as cca
 
 
-def random_part_lens(max_n_parts, max_part_size):
+def random_part_lens(max_n_parts, max_part_size) -> tuple[int, ...]:
     return tuple(random.randint(1, max_part_size) for _ in range(random.randint(1, max_n_parts)))
 
 
@@ -31,7 +31,7 @@ def random_part_lens(max_n_parts, max_part_size):
     'left_part_lens,right_part_lens',
     [tuple(random_part_lens(7, 2) for _ in ('left', 'right')) for _ in range(5)],
 )
-def test_shift_swap_network_gate_acquaintance_opps(left_part_lens, right_part_lens):
+def test_shift_swap_network_gate_acquaintance_opps(left_part_lens, right_part_lens) -> None:
 
     gate = cca.ShiftSwapNetworkGate(left_part_lens, right_part_lens)
     n_qubits = gate.qubit_count()
@@ -45,6 +45,7 @@ def test_shift_swap_network_gate_acquaintance_opps(left_part_lens, right_part_le
     # expected opps
     i = 0
     sides = ('left', 'right')
+    parts: dict[str, list[set[int]]]
     parts = {side: [] for side in sides}
     for side, part_lens in zip(sides, (left_part_lens, right_part_lens)):
         for part_len in part_lens:
@@ -205,7 +206,7 @@ circuit_diagrams = {
 
 
 @pytest.mark.parametrize('left_part_lens,right_part_lens', set(key[1:] for key in circuit_diagrams))
-def test_shift_swap_network_gate_diagrams(left_part_lens, right_part_lens):
+def test_shift_swap_network_gate_diagrams(left_part_lens, right_part_lens) -> None:
 
     gate = cca.ShiftSwapNetworkGate(left_part_lens, right_part_lens)
     n_qubits = gate.qubit_count()
@@ -220,7 +221,7 @@ def test_shift_swap_network_gate_diagrams(left_part_lens, right_part_lens):
     cirq.testing.assert_has_diagram(circuit, diagram)
 
 
-def test_shift_swap_network_gate_bad_part_lens():
+def test_shift_swap_network_gate_bad_part_lens() -> None:
     with pytest.raises(ValueError):
         cca.ShiftSwapNetworkGate((0, 1, 1), (2, 2))
     with pytest.raises(ValueError):
@@ -235,7 +236,7 @@ def test_shift_swap_network_gate_bad_part_lens():
     'left_part_lens,right_part_lens',
     [tuple(random_part_lens(2, 2) for _ in ('left', 'right')) for _ in range(5)],
 )
-def test_shift_swap_network_gate_repr(left_part_lens, right_part_lens):
+def test_shift_swap_network_gate_repr(left_part_lens, right_part_lens) -> None:
     gate = cca.ShiftSwapNetworkGate(left_part_lens, right_part_lens)
     cirq.testing.assert_equivalent_repr(gate)
 
@@ -247,7 +248,7 @@ def test_shift_swap_network_gate_repr(left_part_lens, right_part_lens):
     'left_part_lens,right_part_lens',
     [tuple(random_part_lens(2, 2) for _ in ('left', 'right')) for _ in range(5)],
 )
-def test_shift_swap_network_gate_permutation(left_part_lens, right_part_lens):
+def test_shift_swap_network_gate_permutation(left_part_lens, right_part_lens) -> None:
     gate = cca.ShiftSwapNetworkGate(left_part_lens, right_part_lens)
     n_qubits = gate.qubit_count()
     cca.testing.assert_permutation_decomposition_equivalence(gate, n_qubits)

cirq/contrib/acquaintance/strategies/quartic_paired_test.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import itertools
 import random
+from typing import Iterable
 
 import pytest
 
@@ -23,27 +24,28 @@ import cirq
 import cirq.contrib.acquaintance as cca
 
 
-def test_bad_qubit_pairs():
+def test_bad_qubit_pairs() -> None:
     a, b, c, d, e = cirq.LineQubit.range(5)
     bad_qubit_pairs = [(a, b), (c, d), (e,)]
     with pytest.raises(ValueError):
         cca.strategies.quartic_paired.qubit_pairs_to_qubit_order(bad_qubit_pairs)
 
 
-def random_index_pairs(n_pairs: int):
+def random_index_pairs(n_pairs: int) -> tuple[tuple[int, int], ...]:
     indices = list(range(2 * n_pairs))
     random.shuffle(indices)
-    return tuple(indices[2 * i : 2 * (i + 1)] for i in range(n_pairs))
+    return tuple(zip(indices[0::2], indices[1::2]))
 
 
 @pytest.mark.parametrize(
     'index_pairs', [random_index_pairs(n_pairs) for n_pairs in range(2, 7) for _ in range(2)]
 )
-def test_quartic_paired_acquaintances(index_pairs):
+def test_quartic_paired_acquaintances(index_pairs) -> None:
     n_pairs = len(index_pairs)
-    qubit_pairs = tuple(tuple(cirq.LineQubit(x) for x in index_pair) for index_pair in index_pairs)
+    qubit_pairs: Iterable[tuple[cirq.Qid, cirq.Qid]]
+    qubit_pairs = tuple((cirq.LineQubit(x), cirq.LineQubit(y)) for x, y in index_pairs)
     strategy, qubits = cca.quartic_paired_acquaintance_strategy(qubit_pairs)
-    initial_mapping = {q: q.x for q in qubits}
+    initial_mapping = {q: q.x for q in qubits}  # type: ignore[attr-defined]
     opps = cca.get_logical_acquaintance_opportunities(strategy, initial_mapping)
     assert set(len(opp) for opp in opps) == set([2, 4])
     quadratic_opps = set(opp for opp in opps if len(opp) == 2)

cirq/contrib/bayesian_network/bayesian_network_gate.py

@@ -97,7 +97,7 @@ class BayesianNetworkGate(raw_types.Gate):
     def __init__(
         self,
         init_probs: list[tuple[str, float | None]],
-        arc_probs: list[tuple[str, tuple[str], list[float]]],
+        arc_probs: list[tuple[str, tuple[str, ...], list[float]]],
     ):
         """Builds a BayesianNetworkGate.
 
@@ -195,7 +195,7 @@ class BayesianNetworkGate(raw_types.Gate):
             list[tuple[str, float | None]], [(param, init_prob) for param, init_prob in init_probs]
         )
         converted_cond_probs = cast(
-            list[tuple[str, tuple[str], list[float]]],
+            list[tuple[str, tuple[str, ...], list[float]]],
             [(target, tuple(params), cond_probs) for target, params, cond_probs in arc_probs],
         )
         return cls(converted_init_probs, converted_cond_probs)

cirq/contrib/bayesian_network/bayesian_network_gate_test.py

@@ -21,14 +21,14 @@ import cirq
 import cirq.contrib.bayesian_network as ccb
 
 
-def test_basic_properties():
+def test_basic_properties() -> None:
     gate = ccb.BayesianNetworkGate([('q0', None), ('q1', None), ('q2', None)], [])
 
     assert gate._has_unitary_()
     assert gate._qid_shape_() == (2, 2, 2)
 
 
-def test_incorrect_constructor():
+def test_incorrect_constructor() -> None:
     # Success building.
     ccb.BayesianNetworkGate([('q0', 0.0), ('q1', None)], [('q1', ('q0',), [0.0, 0.0])])
 
@@ -37,7 +37,9 @@ def test_incorrect_constructor():
 
     # This is an easy mistake where the tuple for q0 doesn't have the comma at the end.
     with pytest.raises(ValueError, match='Conditional prob params must be a tuple.'):
-        ccb.BayesianNetworkGate([('q0', 0.0), ('q1', None)], [('q1', ('q0'), [0.0, 0.0])])
+        ccb.BayesianNetworkGate(
+            [('q0', 0.0), ('q1', None)], [('q1', ('q0'), [0.0, 0.0])]  # type: ignore
+        )
 
     with pytest.raises(ValueError, match='Incorrect number of conditional probs.'):
         ccb.BayesianNetworkGate([('q0', 0.0), ('q1', None)], [('q1', ('q0',), [0.0])])
@@ -46,7 +48,7 @@ def test_incorrect_constructor():
         ccb.BayesianNetworkGate([('q0', 0.0), ('q1', None)], [('q1', ('q0',), [2016.0913, 0.0])])
 
 
-def test_repr():
+def test_repr() -> None:
     gate = ccb.BayesianNetworkGate([('q0', 0.0), ('q1', None)], [('q1', ('q0',), [0.0, 0.0])])
 
     assert repr(gate) == (
@@ -56,7 +58,7 @@ def test_repr():
 
 
 @pytest.mark.parametrize('input_prob', [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
-def test_prob_encoding(input_prob):
+def test_prob_encoding(input_prob) -> None:
     q = cirq.NamedQubit('q')
     gate = ccb.BayesianNetworkGate([('q', input_prob)], [])
     circuit = cirq.Circuit(gate.on(q))
@@ -82,7 +84,7 @@ def test_prob_encoding(input_prob):
     ],
 )
 @pytest.mark.parametrize('decompose', [True, False])
-def test_initial_probs(p0, p1, p2, expected_probs, decompose):
+def test_initial_probs(p0, p1, p2, expected_probs, decompose) -> None:
     q0, q1, q2 = cirq.LineQubit.range(3)
     gate = ccb.BayesianNetworkGate([('q0', p0), ('q1', p1), ('q2', p2)], [])
     if decompose:
@@ -102,7 +104,7 @@ def test_initial_probs(p0, p1, p2, expected_probs, decompose):
     [(0.0, 0.0, 0.1), (0.0, 1.0, 0.2), (1.0, 0.0, 0.3), (1.0, 1.0, 0.4)],
 )
 @pytest.mark.parametrize('decompose', [True, False])
-def test_arc_probs(input_prob_q0, input_prob_q1, expected_prob_q2, decompose):
+def test_arc_probs(input_prob_q0, input_prob_q1, expected_prob_q2, decompose) -> None:
     q0, q1, q2 = cirq.LineQubit.range(3)
     gate = ccb.BayesianNetworkGate(
         [('q0', input_prob_q0), ('q1', input_prob_q1), ('q2', None)],
@@ -122,7 +124,7 @@ def test_arc_probs(input_prob_q0, input_prob_q1, expected_prob_q2, decompose):
     np.testing.assert_almost_equal(actual_prob_q2_is_one, expected_prob_q2, decimal=4)
 
 
-def test_repro_figure_10_of_paper():
+def test_repro_figure_10_of_paper() -> None:
     # We try to create the network of figure 10 and check that the probabilities are the same as
     # the ones in table 10 of https://arxiv.org/abs/2004.14803.
     ir = cirq.NamedQubit('q4_IR')

cirq/contrib/circuitdag/circuit_dag.py

@@ -44,7 +44,7 @@ class Unique(Generic[T]):
     def __repr__(self) -> str:
         return f'cirq.contrib.Unique({id(self)}, {self.val!r})'
 
-    def __lt__(self, other):
+    def __lt__(self, other) -> bool:
         if not isinstance(other, type(self)):
             return NotImplemented
         return id(self) < id(other)
@@ -176,7 +176,7 @@ class CircuitDag(networkx.DiGraph):
     def all_operations(self) -> Iterator[cirq.Operation]:
         return (node.val for node in self.ordered_nodes())
 
-    def all_qubits(self):
+    def all_qubits(self) -> frozenset[cirq.Qid]:
         return frozenset(q for node in self.nodes for q in node.val.qubits)
 
     def to_circuit(self) -> cirq.Circuit:

cirq/contrib/circuitdag/circuit_dag_test.py

@@ -27,7 +27,7 @@ class FakeDevice(cirq.Device):
     pass
 
 
-def test_wrapper_eq():
+def test_wrapper_eq() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(cirq.contrib.CircuitDag.make_node(cirq.X(q0)))
@@ -36,7 +36,7 @@ def test_wrapper_eq():
     eq.add_equality_group(cirq.contrib.CircuitDag.make_node(cirq.X(q1)))
 
 
-def test_wrapper_cmp():
+def test_wrapper_cmp() -> None:
     u0 = cirq.contrib.Unique(0)
     u1 = cirq.contrib.Unique(1)
     # The ordering of Unique instances is unpredictable
@@ -51,14 +51,14 @@ def test_wrapper_cmp():
     assert u1 >= u0
 
 
-def test_wrapper_cmp_failure():
+def test_wrapper_cmp_failure() -> None:
     with pytest.raises(TypeError):
         _ = object() < cirq.contrib.Unique(1)
     with pytest.raises(TypeError):
         _ = cirq.contrib.Unique(1) < object()
 
 
-def test_wrapper_repr():
+def test_wrapper_repr() -> None:
     q0 = cirq.LineQubit(0)
 
     node = cirq.contrib.CircuitDag.make_node(cirq.X(q0))
@@ -66,14 +66,14 @@ def test_wrapper_repr():
     assert repr(node) == expected
 
 
-def test_init():
+def test_init() -> None:
     dag = cirq.contrib.CircuitDag()
     assert networkx.dag.is_directed_acyclic_graph(dag)
     assert list(dag.nodes()) == []
     assert list(dag.edges()) == []
 
 
-def test_append():
+def test_append() -> None:
     q0 = cirq.LineQubit(0)
     dag = cirq.contrib.CircuitDag()
     dag.append(cirq.X(q0))
@@ -83,7 +83,7 @@ def test_append():
     assert [(n1.val, n2.val) for n1, n2 in dag.edges()] == [(cirq.X(q0), cirq.Y(q0))]
 
 
-def test_two_identical_ops():
+def test_two_identical_ops() -> None:
     q0 = cirq.LineQubit(0)
     dag = cirq.contrib.CircuitDag()
     dag.append(cirq.X(q0))
@@ -98,7 +98,7 @@ def test_two_identical_ops():
     }
 
 
-def test_from_ops():
+def test_from_ops() -> None:
     q0 = cirq.LineQubit(0)
     dag = cirq.contrib.CircuitDag.from_ops(cirq.X(q0), cirq.Y(q0))
     assert networkx.dag.is_directed_acyclic_graph(dag)
@@ -106,7 +106,7 @@ def test_from_ops():
     assert [(n1.val, n2.val) for n1, n2 in dag.edges()] == [(cirq.X(q0), cirq.Y(q0))]
 
 
-def test_from_circuit():
+def test_from_circuit() -> None:
     q0 = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.X(q0), cirq.Y(q0))
     dag = cirq.contrib.CircuitDag.from_circuit(circuit)
@@ -116,14 +116,14 @@ def test_from_circuit():
     assert sorted(circuit.all_qubits()) == sorted(dag.all_qubits())
 
 
-def test_to_empty_circuit():
+def test_to_empty_circuit() -> None:
     circuit = cirq.Circuit()
     dag = cirq.contrib.CircuitDag.from_circuit(circuit)
     assert networkx.dag.is_directed_acyclic_graph(dag)
     assert circuit == dag.to_circuit()
 
 
-def test_to_circuit():
+def test_to_circuit() -> None:
     q0 = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.X(q0), cirq.Y(q0))
     dag = cirq.contrib.CircuitDag.from_circuit(circuit)
@@ -137,7 +137,7 @@ def test_to_circuit():
     )
 
 
-def test_equality():
+def test_equality() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     circuit1 = cirq.Circuit(
         cirq.X(q0), cirq.Y(q0), cirq.Z(q1), cirq.CZ(q0, q1), cirq.X(q1), cirq.Y(q1), cirq.Z(q0)
@@ -167,7 +167,7 @@ def test_equality():
     eq.add_equality_group(cirq.contrib.CircuitDag.from_circuit(circuit4))
 
 
-def test_larger_circuit():
+def test_larger_circuit() -> None:
     q0, q1, q2, q3 = [
         cirq.GridQubit(0, 5),
         cirq.GridQubit(1, 5),
@@ -213,7 +213,7 @@ def test_larger_circuit():
 
 
 @pytest.mark.parametrize('circuit', [cirq.testing.random_circuit(10, 10, 0.5) for _ in range(3)])
-def test_is_maximalist(circuit):
+def test_is_maximalist(circuit) -> None:
     dag = cirq.contrib.CircuitDag.from_circuit(circuit)
     transitive_closure = networkx.dag.transitive_closure(dag)
     assert cirq.contrib.CircuitDag(incoming_graph_data=transitive_closure) == dag
@@ -232,7 +232,7 @@ def _get_circuits_and_is_blockers():
 
 
 @pytest.mark.parametrize('circuit, is_blocker', _get_circuits_and_is_blockers())
-def test_findall_nodes_until_blocked(circuit, is_blocker):
+def test_findall_nodes_until_blocked(circuit, is_blocker) -> None:
     dag = cirq.contrib.CircuitDag.from_circuit(circuit)
     all_nodes = list(dag.ordered_nodes())
     found_nodes = list(dag.findall_nodes_until_blocked(is_blocker))