cirq-core 1.5.0.dev20250409222543__py3-none-any.whl → 1.6.0__py3-none-any.whl

cirq/contrib/acquaintance/bipartite_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import itertools
 
 import pytest
@@ -162,9 +164,11 @@ circuit_diagrams = {
                 │   │         │   │
 3: ─────────────█───1↦0───────█───1↦0─────────────
 """,
-    ('decomposed', cca.BipartiteGraphType.COMPLETE, 3):
-    # pylint: disable=line-too-long
-    """
+    (
+        'decomposed',
+        cca.BipartiteGraphType.COMPLETE,
+        3,
+    ): """
 0: ───────────────────────█───0↦1───────────────────────────█───0↦1───────────────────────
                           │   │                             │   │
 1: ─────────────█───0↦1───█───1↦0───█───0↦1───────█───0↦1───█───1↦0───█───0↦1─────────────
@@ -199,8 +203,7 @@ circuit_diagrams = {
                                     │   │                                                 │   │
 7: ─────────────────────────────────█───1↦0───────────────────────────────────────────────█───1↦0─────────────────────────────────
 
-""",
-    # pylint: enable=line-too-long
+""",  # noqa: E501
     (
         'decomposed',
         cca.BipartiteGraphType.MATCHING,
@@ -267,7 +270,7 @@ circuit_diagrams = {
 @pytest.mark.parametrize(
     'subgraph,part_size', itertools.product(cca.BipartiteGraphType, range(1, 5))
 )
-def test_circuit_diagrams(part_size, subgraph):
+def test_circuit_diagrams(part_size, subgraph) -> None:
     qubits = cirq.LineQubit.range(2 * part_size)
     gate = cca.BipartiteSwapNetworkGate(subgraph, part_size)
     circuit = cirq.Circuit(gate(*qubits))
@@ -282,10 +285,10 @@ def test_circuit_diagrams(part_size, subgraph):
     cirq.testing.assert_has_diagram(circuit, diagram)
 
 
-def test_bad_args():
+def test_bad_args() -> None:
     gate = cca.BipartiteSwapNetworkGate(cca.BipartiteGraphType.COMPLETE, 2)
     qubits = cirq.LineQubit.range(4)
-    gate.subgraph = 'not a subgraph'
+    gate.subgraph = 'not a subgraph'  # type: ignore[assignment]
     args = cirq.CircuitDiagramInfoArgs(
         known_qubits=None,
         known_qubit_count=None,
@@ -303,7 +306,7 @@ def test_bad_args():
     with pytest.raises(ValueError):
         gate._decompose_(qubits[:3])
 
-    gate.subgraph = 'unimplemented subgraph'
+    gate.subgraph = 'unimplemented subgraph'  # type: ignore[assignment]
     with pytest.raises(NotImplementedError):
         gate._decompose_(qubits)
 
@@ -312,7 +315,7 @@ def test_bad_args():
         gate._circuit_diagram_info_(args)
 
 
-def test_bipartite_swap_network_acquaintance_size():
+def test_bipartite_swap_network_acquaintance_size() -> None:
     qubits = cirq.LineQubit.range(4)
     gate = cca.BipartiteSwapNetworkGate(cca.BipartiteGraphType.COMPLETE, 2)
     assert cca.get_acquaintance_size(gate(*qubits)) == 2
@@ -321,7 +324,7 @@ def test_bipartite_swap_network_acquaintance_size():
 @pytest.mark.parametrize(
     'subgraph,part_size', itertools.product(cca.BipartiteGraphType, range(1, 3))
 )
-def test_repr(subgraph, part_size):
+def test_repr(subgraph, part_size) -> None:
     gate = cca.BipartiteSwapNetworkGate(subgraph, part_size)
     cirq.testing.assert_equivalent_repr(gate)
 
@@ -332,18 +335,20 @@ def test_repr(subgraph, part_size):
 @pytest.mark.parametrize(
     'subgraph,part_size', itertools.product(cca.BipartiteGraphType, range(1, 6))
 )
-def test_decomposition_permutation_consistency(part_size, subgraph):
+def test_decomposition_permutation_consistency(part_size, subgraph) -> None:
     gate = cca.BipartiteSwapNetworkGate(subgraph, part_size)
     qubits = cirq.LineQubit.range(2 * part_size)
+    mapping: dict[cirq.Qid, int]
     mapping = {q: i for i, q in enumerate(qubits)}
     cca.update_mapping(mapping, gate._decompose_(qubits))
     permutation = gate.permutation()
     assert {qubits[i]: j for i, j in permutation.items()} == mapping
 
 
-def test_bad_number_of_qubits():
+def test_bad_number_of_qubits() -> None:
     gate = cca.BipartiteSwapNetworkGate(cca.BipartiteGraphType.COMPLETE, 6)
     qubits = cirq.LineQubit.range(6)
+    mapping: dict[cirq.Qid, int]
     mapping = {q: i for i, q in enumerate(qubits)}
     with pytest.raises(ValueError, match='len'):
         cca.update_mapping(mapping, gate._decompose_(qubits))

cirq/contrib/acquaintance/devices.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from cirq import circuits, devices, ops
 from cirq.contrib.acquaintance.bipartite import BipartiteSwapNetworkGate
@@ -42,7 +42,7 @@ class AcquaintanceDevice(devices.Device, metaclass=abc.ABCMeta):
             )
 
 
-def get_acquaintance_size(obj: Union[circuits.Circuit, ops.Operation]) -> int:
+def get_acquaintance_size(obj: circuits.Circuit | ops.Operation) -> int:
     """The maximum number of qubits to be acquainted with each other."""
     if isinstance(obj, circuits.Circuit):
         return max(tuple(get_acquaintance_size(op) for op in obj.all_operations()) or (0,))

cirq/contrib/acquaintance/devices_test.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import pytest
 
 import cirq
 import cirq.contrib.acquaintance as cca
 
 
-def test_acquaintance_device():
+def test_acquaintance_device() -> None:
     with pytest.raises(ValueError):
         op = cirq.X(cirq.NamedQubit('q'))
         cca.UnconstrainedAcquaintanceDevice.validate_operation(op)
@@ -29,6 +31,6 @@ def test_acquaintance_device():
     cca.UnconstrainedAcquaintanceDevice.validate_operation(swap_network(*qubits))
 
 
-def test_not_operation():
+def test_not_operation() -> None:
     with pytest.raises(TypeError):
-        _ = cca.get_acquaintance_size(cirq.LineQubit(1))
+        _ = cca.get_acquaintance_size(cirq.LineQubit(1))  # type: ignore[arg-type]

cirq/contrib/acquaintance/executor.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import abc
 from collections import defaultdict
-from typing import DefaultDict, Dict, Iterator, Optional, Sequence, TYPE_CHECKING
+from typing import Iterator, Sequence, TYPE_CHECKING
 
 from cirq import circuits, devices, ops, protocols, transformers
 from cirq.contrib.acquaintance.gates import AcquaintanceOpportunityGate
@@ -99,7 +99,7 @@ class StrategyExecutorTransformer:
         self._mapping = execution_strategy.initial_mapping.copy()
 
     def __call__(
-        self, circuit: circuits.AbstractCircuit, context: Optional[cirq.TransformerContext] = None
+        self, circuit: circuits.AbstractCircuit, context: cirq.TransformerContext | None = None
     ) -> circuits.Circuit:
         """Executes an acquaintance strategy using cirq.map_operations_and_unroll and
         mutates initial mapping.
@@ -174,7 +174,7 @@ class GreedyExecutionStrategy(ExecutionStrategy):
         self,
         gates: LogicalGates,
         initial_mapping: LogicalMapping,
-        device: Optional[cirq.Device] = None,
+        device: cirq.Device | None = None,
     ) -> None:
         """Inits GreedyExecutionStrategy.
 
@@ -214,13 +214,13 @@ class GreedyExecutionStrategy(ExecutionStrategy):
                 yield gate(*[index_to_qubit[i] for i in gate_indices])
 
     @staticmethod
-    def canonicalize_gates(gates: LogicalGates) -> Dict[frozenset, LogicalGates]:
+    def canonicalize_gates(gates: LogicalGates) -> dict[frozenset, LogicalGates]:
         """Canonicalizes a set of gates by the qubits they act on.
 
         Takes a set of gates specified by ordered sequences of logical
         indices, and groups those that act on the same qubits regardless of
         order."""
-        canonicalized_gates: DefaultDict[frozenset, LogicalGates] = defaultdict(dict)
+        canonicalized_gates: defaultdict[frozenset, LogicalGates] = defaultdict(dict)
         for indices, gate in gates.items():
             indices = tuple(indices)
             canonicalized_gates[frozenset(indices)][indices] = gate

cirq/contrib/acquaintance/executor_test.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from itertools import combinations
 from string import ascii_lowercase
-from typing import Dict, Sequence, Tuple
+from typing import Any, Sequence
 
 import numpy as np
 import pytest
@@ -36,11 +38,12 @@ class ExampleGate(cirq.Gate):
         return self._wire_symbols
 
 
-def test_executor_explicit():
+def test_executor_explicit() -> None:
     num_qubits = 8
     qubits = cirq.LineQubit.range(num_qubits)
     circuit = cca.complete_acquaintance_strategy(qubits, 2)
 
+    gates: dict[tuple[Any, ...], cirq.Gate]
     gates = {
         (i, j): ExampleGate([str(k) for k in ij])
         for ij in combinations(range(num_qubits), 2)
@@ -51,10 +54,11 @@ def test_executor_explicit():
     assert execution_strategy.device == cirq.UNCONSTRAINED_DEVICE
 
     with pytest.raises(ValueError):
-        executor = cca.StrategyExecutorTransformer(None)
+        executor = cca.StrategyExecutorTransformer(None)  # type: ignore[arg-type]
     executor = cca.StrategyExecutorTransformer(execution_strategy)
 
     with pytest.raises(NotImplementedError):
+        bad_gates: dict[tuple[Any, ...], cirq.Gate]
         bad_gates = {(0,): ExampleGate(['0']), (0, 1): ExampleGate(['0', '1'])}
         cca.GreedyExecutionStrategy(bad_gates, initial_mapping)
 
@@ -79,13 +83,13 @@ def test_executor_explicit():
 6: ───6───7───╲0╱───7───4───╱1╲───4───6───╲0╱───6───2───╱1╲───2───4───╲0╱───4───0───╱1╲───0───2───╲0╱───2───1───╱1╲───
       │   │   │                   │   │   │                   │   │   │                   │   │   │
 7: ───7───6───╱1╲─────────────────6───4───╱1╲─────────────────4───2───╱1╲─────────────────2───0───╱1╲─────────────────
-    """.strip()
+    """.strip()  # noqa: E501
     ct.assert_has_diagram(circuit, expected_text_diagram)
 
 
 def random_diagonal_gates(
     num_qubits: int, acquaintance_size: int
-) -> Dict[Tuple[cirq.Qid, ...], cirq.Gate]:
+) -> dict[tuple[cirq.Qid, ...], cirq.Gate]:
 
     return {
         Q: cirq.DiagonalGate(np.random.random(2**acquaintance_size))
@@ -104,8 +108,8 @@ def random_diagonal_gates(
     ],
 )
 def test_executor_random(
-    num_qubits: int, acquaintance_size: int, gates: Dict[Tuple[cirq.Qid, ...], cirq.Gate]
-):
+    num_qubits: int, acquaintance_size: int, gates: dict[tuple[cirq.Qid, ...], cirq.Gate]
+) -> None:
     qubits = cirq.LineQubit.range(num_qubits)
     circuit = cca.complete_acquaintance_strategy(qubits, acquaintance_size)
 
@@ -124,7 +128,7 @@ def test_executor_random(
     np.testing.assert_allclose(actual=actual_unitary, desired=expected_unitary, verbose=True)
 
 
-def test_acquaintance_operation():
+def test_acquaintance_operation() -> None:
     n = 5
     physical_qubits = tuple(cirq.LineQubit.range(n))
     logical_qubits = tuple(cirq.NamedQubit(s) for s in ascii_lowercase[:n])
@@ -132,7 +136,7 @@ def test_acquaintance_operation():
     with pytest.raises(ValueError):
         cca.AcquaintanceOperation(physical_qubits[:3], int_indices[:4])
     for logical_indices in (logical_qubits, int_indices):
-        op = cca.AcquaintanceOperation(physical_qubits, logical_indices)
+        op = cca.AcquaintanceOperation(physical_qubits, logical_indices)  # type: ignore[type-var]
         assert op.logical_indices == logical_indices
         assert op.qubits == physical_qubits
         wire_symbols = tuple(f'({i})' for i in logical_indices)

cirq/contrib/acquaintance/gates.py

@@ -18,17 +18,7 @@ import functools
 import itertools
 import math
 import operator
-from typing import (
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    NamedTuple,
-    Optional,
-    Sequence,
-    Tuple,
-    TYPE_CHECKING,
-)
+from typing import Iterable, Iterator, NamedTuple, Sequence, TYPE_CHECKING
 
 from cirq import ops, protocols, value
 from cirq.contrib.acquaintance.permutation import (
@@ -44,14 +34,14 @@ if TYPE_CHECKING:
 
 def operations_to_part_lens(
     qubit_order: Sequence[cirq.Qid], op_tree: cirq.OP_TREE
-) -> Tuple[int, ...]:
+) -> tuple[int, ...]:
     qubit_sort_key = functools.partial(operator.indexOf, qubit_order)
     op_parts = [tuple(sorted(op.qubits, key=qubit_sort_key)) for op in ops.flatten_op_tree(op_tree)]
-    singletons: List[Tuple[cirq.Qid, ...]] = [(q,) for q in set(qubit_order).difference(*op_parts)]
+    singletons: list[tuple[cirq.Qid, ...]] = [(q,) for q in set(qubit_order).difference(*op_parts)]
     part_sort_key = lambda p: min(qubit_sort_key(q) for q in p)
     parts = tuple(tuple(part) for part in sorted(singletons + op_parts, key=part_sort_key))
 
-    if sum(parts, ()) != tuple(qubit_order):
+    if tuple(itertools.chain.from_iterable(parts)) != tuple(qubit_order):
         raise ValueError('sum(parts, ()) != tuple(qubit_order)')
 
     return tuple(len(part) for part in parts)
@@ -87,16 +77,16 @@ def acquaint(*qubits) -> cirq.Operation:
 Layers = NamedTuple(
     'Layers',
     [
-        ('prior_interstitial', List['cirq.Operation']),
-        ('pre', List['cirq.Operation']),
-        ('intra', List['cirq.Operation']),
-        ('post', List['cirq.Operation']),
-        ('posterior_interstitial', List['cirq.Operation']),
+        ('prior_interstitial', list['cirq.Operation']),
+        ('pre', list['cirq.Operation']),
+        ('intra', list['cirq.Operation']),
+        ('post', list['cirq.Operation']),
+        ('posterior_interstitial', list['cirq.Operation']),
     ],
 )
 
 
-def new_layers(**kwargs: List[cirq.Operation]) -> Layers:
+def new_layers(**kwargs: list[cirq.Operation]) -> Layers:
     return Layers._make(kwargs.get(field, []) for field in Layers._fields)
 
 
@@ -106,7 +96,7 @@ def acquaint_insides(
     qubits: Sequence[cirq.Qid],
     before: bool,
     layers: Layers,
-    mapping: Dict[ops.Qid, int],
+    mapping: dict[ops.Qid, int],
 ) -> None:
     """Acquaints each of the qubits with another set specified by an
     acquaintance gate.
@@ -153,12 +143,12 @@ def _get_max_reach(size: int, round_up: bool = True) -> int:
 
 
 def acquaint_and_shift(
-    parts: Tuple[List[cirq.Qid], List[cirq.Qid]],
+    parts: tuple[list[cirq.Qid], list[cirq.Qid]],
     layers: Layers,
-    acquaintance_size: Optional[int],
+    acquaintance_size: int | None,
     swap_gate: cirq.Gate,
-    mapping: Dict[ops.Qid, int],
-):
+    mapping: dict[ops.Qid, int],
+) -> None:
     """Acquaints and shifts a pair of lists of qubits. The first part is
     acquainted with every qubit individually in the second part, and vice
     versa. Operations are grouped into several layers:
@@ -276,7 +266,7 @@ class SwapNetworkGate(PermutationGate):
     def __init__(
         self,
         part_lens: Sequence[int],
-        acquaintance_size: Optional[int] = 0,
+        acquaintance_size: int | None = 0,
         swap_gate: cirq.Gate = ops.SWAP,
     ) -> None:
         super().__init__(sum(part_lens), swap_gate)
@@ -346,13 +336,13 @@ class SwapNetworkGate(PermutationGate):
     def from_operations(
         qubit_order: Sequence[cirq.Qid],
         operations: Sequence[cirq.Operation],
-        acquaintance_size: Optional[int] = 0,
+        acquaintance_size: int | None = 0,
         swap_gate: cirq.Gate = ops.SWAP,
     ) -> SwapNetworkGate:
         part_sizes = operations_to_part_lens(qubit_order, operations)
         return SwapNetworkGate(part_sizes, acquaintance_size, swap_gate)
 
-    def permutation(self) -> Dict[int, int]:
+    def permutation(self) -> dict[int, int]:
         return {i: j for i, j in enumerate(reversed(range(sum(self.part_lens))))}
 
     def __repr__(self) -> str:

cirq/contrib/acquaintance/gates_test.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from itertools import combinations, product
 from random import randint
 from string import ascii_lowercase as alphabet
-from typing import Optional, Sequence, Tuple
+from typing import Sequence
 
 import numpy
 import pytest
@@ -25,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'
@@ -33,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()
@@ -47,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
@@ -123,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:
@@ -139,7 +141,7 @@ def test_acquaint_part_pairs(part_lens):
     assert expected_opps == actual_opps
 
 
-acquaintance_sizes: Tuple[Optional[int], ...] = (None,)
+acquaintance_sizes: tuple[int | None, ...] = (None,)
 acquaintance_sizes += tuple(range(5))
 
 
@@ -152,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)
@@ -196,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)
@@ -221,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):
@@ -235,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)
 
@@ -244,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)
 
@@ -254,20 +256,22 @@ class OtherOperation(cirq.Operation):
         self._qubits = tuple(qubits)
 
     @property
-    def qubits(self) -> Tuple[cirq.Qid, ...]:
+    def qubits(self) -> tuple[cirq.Qid, ...]:
         return self._qubits
 
-    def with_qubits(self, *new_qubits: cirq.Qid) -> 'OtherOperation':
-        return type(self)(self._qubits)
+    def with_qubits(self, *new_qubits: cirq.Qid) -> OtherOperation:
+        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):
@@ -299,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)
@@ -317,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

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import FrozenSet, Iterator, Sequence, Set, TYPE_CHECKING
+from typing import Iterator, Sequence, TYPE_CHECKING
 
 from cirq import devices
 from cirq.contrib import circuitdag
@@ -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/inspection_utils_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from itertools import combinations, product
 
 import pytest
@@ -22,7 +24,7 @@ import cirq.contrib.acquaintance.inspection_utils as inspection_utils
 
 
 @pytest.mark.parametrize('n_qubits, acquaintance_size', product(range(2, 6), range(2, 5)))
-def test_get_logical_acquaintance_opportunities(n_qubits, acquaintance_size):
+def test_get_logical_acquaintance_opportunities(n_qubits, acquaintance_size) -> None:
     qubits = cirq.LineQubit.range(n_qubits)
     acquaintance_strategy = cca.complete_acquaintance_strategy(qubits, acquaintance_size)
     initial_mapping = {q: i for i, q in enumerate(qubits)}
@@ -30,5 +32,5 @@ def test_get_logical_acquaintance_opportunities(n_qubits, acquaintance_size):
     assert opps == set(frozenset(s) for s in combinations(range(n_qubits), acquaintance_size))
 
 
-def test_device():
+def test_device() -> None:
     assert inspection_utils.LogicalAnnotator({}).device == cirq.UNCONSTRAINED_DEVICE

cirq/contrib/acquaintance/mutation_utils.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import collections
-from typing import cast, Dict, List, Optional, Sequence, TYPE_CHECKING, Union
+from typing import cast, Sequence, TYPE_CHECKING
 
 from cirq import circuits, ops, transformers
 from cirq.contrib.acquaintance.devices import get_acquaintance_size
@@ -25,7 +25,7 @@ from cirq.contrib.acquaintance.permutation import PermutationGate
 if TYPE_CHECKING:
     import cirq
 
-STRATEGY_GATE = Union[AcquaintanceOpportunityGate, PermutationGate]
+STRATEGY_GATE = AcquaintanceOpportunityGate | PermutationGate
 
 
 def rectify_acquaintance_strategy(circuit: cirq.Circuit, acquaint_first: bool = True) -> None:
@@ -43,7 +43,7 @@ def rectify_acquaintance_strategy(circuit: cirq.Circuit, acquaint_first: bool =
     """
     rectified_moments = []
     for moment in circuit:
-        gate_type_to_ops: Dict[bool, List[ops.GateOperation]] = collections.defaultdict(list)
+        gate_type_to_ops: dict[bool, list[ops.GateOperation]] = collections.defaultdict(list)
         for op in moment.operations:
             gate_op = cast(ops.GateOperation, op)
             is_acquaintance = isinstance(gate_op.gate, AcquaintanceOpportunityGate)
@@ -59,7 +59,7 @@ def rectify_acquaintance_strategy(circuit: cirq.Circuit, acquaint_first: bool =
 def replace_acquaintance_with_swap_network(
     circuit: cirq.Circuit,
     qubit_order: Sequence[cirq.Qid],
-    acquaintance_size: Optional[int] = 0,
+    acquaintance_size: int | None = 0,
     swap_gate: cirq.Gate = ops.SWAP,
 ) -> bool:
     """Replace every rectified moment with acquaintance gates with a generalized swap network.

cirq/contrib/acquaintance/mutation_utils_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from string import ascii_lowercase as alphabet
 
 import pytest
@@ -20,7 +22,7 @@ import cirq
 import cirq.contrib.acquaintance as cca
 
 
-def test_complete_acquaintance_strategy():
+def test_complete_acquaintance_strategy() -> None:
     qubits = [cirq.NamedQubit(s) for s in alphabet]
 
     with pytest.raises(ValueError):
@@ -136,7 +138,7 @@ a      b      c      d
     assert cca.get_acquaintance_size(cubic_strategy) == 3
 
 
-def test_rectification():
+def test_rectification() -> None:
     qubits = cirq.LineQubit.range(4)
 
     perm_gate = cca.SwapPermutationGate()