Update type annotations

cirq-core/cirq/circuits/circuit_operation_test.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import unittest.mock as mock
-from typing import Optional
 
 import numpy as np
 import pytest
@@ -294,7 +293,7 @@ def test_repeat(add_measurements: bool, use_default_ids_for_initial_rep: bool) -
             _ = op_base.repeat(initial_repetitions)
         initial_repetitions = abs(initial_repetitions)
 
-    op_with_reps: Optional[cirq.CircuitOperation] = None
+    op_with_reps: cirq.CircuitOperation | None = None
     rep_ids = []
     if use_default_ids_for_initial_rep:
         rep_ids = ['0', '1', '2']

cirq-core/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 @@ def validate_operation(self, operation: cirq.Operation) -> None:
             )
 
 
-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-core/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
@@ -66,7 +66,7 @@ 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[LogicalIndex]]:
     acquaintance_dag = get_acquaintance_dag(strategy, initial_mapping)
     logical_acquaintance_opportunities = set()
     for op in acquaintance_dag.all_operations():

cirq-core/cirq/contrib/qasm_import/_lexer.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import re
-from typing import Optional
 
 import ply.lex as lex
 
@@ -117,5 +116,5 @@ def t_error(self, t):
     def input(self, qasm):
         self.lex.input(qasm)
 
-    def token(self) -> Optional[lex.Token]:
+    def token(self) -> lex.Token | None:
         return self.lex.token()

cirq-core/cirq/contrib/routing/router.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Optional, TYPE_CHECKING
+from typing import Callable, TYPE_CHECKING
 
 from cirq import circuits, protocols
 from cirq.contrib.routing.greedy import route_circuit_greedily
@@ -31,8 +31,8 @@ def route_circuit(
     circuit: circuits.Circuit,
     device_graph: nx.Graph,
     *,
-    algo_name: Optional[str] = None,
-    router: Optional[Callable[..., SwapNetwork]] = None,
+    algo_name: str | None = None,
+    router: Callable[..., SwapNetwork] | None = None,
     **kwargs,
 ) -> SwapNetwork:
     """Routes a circuit on a given device.

cirq-core/cirq/devices/device.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import abc
-from typing import FrozenSet, Iterable, Optional, TYPE_CHECKING
+from typing import Iterable, TYPE_CHECKING
 
 import networkx as nx
 
@@ -59,7 +59,7 @@ class Device(metaclass=abc.ABCMeta):
     """
 
     @property
-    def metadata(self) -> Optional[DeviceMetadata]:
+    def metadata(self) -> DeviceMetadata | None:
         """Returns the associated Metadata with the device if applicable.
 
         Returns:
@@ -116,11 +116,11 @@ def __init__(self, qubits: Iterable[cirq.Qid], nx_graph: nx.Graph):
                 directional coupling, undirected edges indicate bi-directional
                 coupling.
         """
-        self._qubits_set: FrozenSet[cirq.Qid] = frozenset(qubits)
+        self._qubits_set: frozenset[cirq.Qid] = frozenset(qubits)
         self._nx_graph = nx_graph
 
     @property
-    def qubit_set(self) -> FrozenSet[cirq.Qid]:
+    def qubit_set(self) -> frozenset[cirq.Qid]:
         """Returns the set of qubits on the device.
 
         Returns:

cirq-core/cirq/experiments/fidelity_estimation.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Mapping, Optional, Sequence, TYPE_CHECKING
+from typing import Callable, Mapping, Sequence, TYPE_CHECKING
 
 import numpy as np
 
@@ -141,7 +141,7 @@ def xeb_fidelity(
     circuit: cirq.Circuit,
     bitstrings: Sequence[int],
     qubit_order: QubitOrderOrList = QubitOrder.DEFAULT,
-    amplitudes: Optional[Mapping[int, complex]] = None,
+    amplitudes: Mapping[int, complex] | None = None,
     estimator: Callable[[int, Sequence[float]], float] = linear_xeb_fidelity_from_probabilities,
 ) -> float:
     """Estimates XEB fidelity from one circuit using user-supplied estimator.
@@ -206,7 +206,7 @@ def linear_xeb_fidelity(
     circuit: cirq.Circuit,
     bitstrings: Sequence[int],
     qubit_order: QubitOrderOrList = QubitOrder.DEFAULT,
-    amplitudes: Optional[Mapping[int, complex]] = None,
+    amplitudes: Mapping[int, complex] | None = None,
 ) -> float:
     """Estimates XEB fidelity from one circuit using linear estimator."""
     return xeb_fidelity(
@@ -222,7 +222,7 @@ def log_xeb_fidelity(
     circuit: cirq.Circuit,
     bitstrings: Sequence[int],
     qubit_order: QubitOrderOrList = QubitOrder.DEFAULT,
-    amplitudes: Optional[Mapping[int, complex]] = None,
+    amplitudes: Mapping[int, complex] | None = None,
 ) -> float:
     """Estimates XEB fidelity from one circuit using logarithmic estimator."""
     return xeb_fidelity(

cirq-core/cirq/experiments/t1_decay_experiment.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, cast, Optional, Sequence, TYPE_CHECKING
+from typing import Any, cast, Sequence, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -146,7 +146,7 @@ def exp_decay(x, t1, a, b):
             return np.nan
 
     def plot(
-        self, ax: Optional[plt.Axes] = None, include_fit: bool = False, **plot_kwargs: Any
+        self, ax: plt.Axes | None = None, include_fit: bool = False, **plot_kwargs: Any
     ) -> plt.Axes:
         """Plots the excited state probability vs the amount of delay.
 

cirq-core/cirq/experiments/xeb_fitting.py

@@ -506,7 +506,7 @@ def _mean_infidelity(angles):
     )
     return XEBCharacterizationResult(
         optimization_results={pair: optimization_result},
-        final_params={pair: final_params},  # type: ignore[dict-item]
+        final_params={pair: final_params},
         fidelities_df=fidelities_df,
     )
 

cirq-core/cirq/interop/quirk/cells/measurement_cells.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import cast, Iterable, Iterator, Optional, TYPE_CHECKING
+from typing import cast, Iterable, Iterator, TYPE_CHECKING
 
 from cirq import ops
 from cirq.interop.quirk.cells.cell import CellMaker, ExplicitOperationsCell
@@ -30,7 +30,7 @@ def generate_all_measurement_cell_makers() -> Iterator[CellMaker]:
     yield _measurement("XDetector", basis_change=ops.Y**0.5)
 
 
-def _measurement(identifier: str, basis_change: Optional[cirq.Gate] = None) -> CellMaker:
+def _measurement(identifier: str, basis_change: cirq.Gate | None = None) -> CellMaker:
     return CellMaker(
         identifier=identifier,
         size=1,

cirq-core/cirq/linalg/combinators.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import functools
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -27,7 +27,7 @@
     from numpy.typing import ArrayLike, DTypeLike
 
 
-def kron(*factors: Union[np.ndarray, complex], shape_len: int = 2) -> np.ndarray:
+def kron(*factors: np.ndarray | complex, shape_len: int = 2) -> np.ndarray:
     """Computes the kronecker product of a sequence of values.
 
     A *args version of lambda args: functools.reduce(np.kron, args).
@@ -58,7 +58,7 @@ def kron(*factors: Union[np.ndarray, complex], shape_len: int = 2) -> np.ndarray
 )
 
 
-def kron_with_controls(*factors: Union[np.ndarray, complex]) -> np.ndarray:
+def kron_with_controls(*factors: np.ndarray | complex) -> np.ndarray:
     """Computes the kronecker product of a sequence of values and control tags.
 
     Use `cirq.CONTROL_TAG` to represent controls. Any entry of the output

cirq-core/cirq/linalg/predicates.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from types import EllipsisType
-from typing import cast, Sequence, Union
+from typing import cast, Sequence
 
 import numpy as np
 
@@ -302,7 +302,7 @@ def slice_for_qubits_equal_to(
     out_size = (
         cast(int, num_qubits) if out_size_specified else max(target_qubit_axes, default=-1) + 1
     )
-    result = cast(list[Union[slice, int, EllipsisType]], [slice(None)] * out_size)
+    result = cast(list[slice | int | EllipsisType], [slice(None)] * out_size)
     if not out_size_specified:
         result.append(Ellipsis)
     if qid_shape is None:

cirq-core/cirq/linalg/tolerance.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable, TYPE_CHECKING, Union
+from typing import Iterable, TYPE_CHECKING
 
 import numpy as np
 
@@ -36,7 +36,7 @@ def all_near_zero(a: ArrayLike, *, atol: float = 1e-8) -> bool:
 
 
 def all_near_zero_mod(
-    a: Union[float, Iterable[float], np.ndarray], period: float, *, atol: float = 1e-8
+    a: float | Iterable[float] | np.ndarray, period: float, *, atol: float = 1e-8
 ) -> bool:
     """Checks if the tensor's elements are all near multiples of the period.
 

cirq-core/cirq/ops/arithmetic_operation_test.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Sequence, Union
+from typing import Sequence
 
 import numpy as np
 import pytest
@@ -89,9 +89,7 @@ def test_the_tests() -> None:
 def test_arithmetic_gate_apply_unitary() -> None:
     class Add(cirq.ArithmeticGate):
         def __init__(
-            self,
-            target_register: Union[int, Sequence[int]],
-            input_register: Union[int, Sequence[int]],
+            self, target_register: int | Sequence[int], input_register: int | Sequence[int]
         ):
             self.target_register = target_register
             self.input_register = input_register

cirq-core/cirq/ops/common_gate_families.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Any, cast, Optional, Type, Union
+from typing import Any, cast
 
 from cirq import protocols
 from cirq.ops import eigen_gate, gateset, parallel_gate, raw_types
@@ -25,7 +25,7 @@
 class AnyUnitaryGateFamily(gateset.GateFamily):
     """GateFamily which accepts any N-Qubit unitary gate."""
 
-    def __init__(self, num_qubits: Optional[int] = None) -> None:
+    def __init__(self, num_qubits: int | None = None) -> None:
         """Init AnyUnitaryGateFamily
 
         Args:
@@ -65,7 +65,7 @@ def _from_json_dict_(cls, num_qubits, **kwargs):
 class AnyIntegerPowerGateFamily(gateset.GateFamily):
     """GateFamily which accepts instances of a given `cirq.EigenGate`, raised to integer power."""
 
-    def __init__(self, gate: Type[eigen_gate.EigenGate]) -> None:
+    def __init__(self, gate: type[eigen_gate.EigenGate]) -> None:
         """Init AnyIntegerPowerGateFamily
 
         Args:
@@ -129,11 +129,11 @@ class ParallelGateFamily(gateset.GateFamily):
 
     def __init__(
         self,
-        gate: Union[Type[raw_types.Gate], raw_types.Gate],
+        gate: type[raw_types.Gate] | raw_types.Gate,
         *,
-        name: Optional[str] = None,
-        description: Optional[str] = None,
-        max_parallel_allowed: Optional[int] = None,
+        name: str | None = None,
+        description: str | None = None,
+        max_parallel_allowed: int | None = None,
     ) -> None:
         """Inits ParallelGateFamily
 

cirq-core/cirq/ops/controlled_gate_test.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from types import EllipsisType, NotImplementedType
-from typing import Any, cast, Sequence, Union
+from typing import Any, cast, Sequence
 
 import numpy as np
 import pytest
@@ -47,7 +47,7 @@ def __init__(self):
     def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray | NotImplementedType:
         assert len(args.axes) == 1
         a = args.axes[0]
-        seed = cast(tuple[Union[int, slice, EllipsisType], ...], (slice(None),))
+        seed = cast(tuple[int | slice | EllipsisType, ...], (slice(None),))
         zero = seed * a + (0, Ellipsis)
         one = seed * a + (1, Ellipsis)
         result = np.zeros(args.target_tensor.shape, args.target_tensor.dtype)

cirq-core/cirq/ops/controlled_operation_test.py

@@ -17,7 +17,7 @@
 import itertools
 import re
 from types import EllipsisType, NotImplementedType
-from typing import cast, Union
+from typing import cast
 
 import numpy as np
 import pytest
@@ -50,7 +50,7 @@ def __init__(self):
     def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray | NotImplementedType:
         assert len(args.axes) == 1
         a = args.axes[0]
-        seed = cast(tuple[Union[int, slice, EllipsisType], ...], (slice(None),))
+        seed = cast(tuple[int | slice | EllipsisType, ...], (slice(None),))
         zero = seed * a + (0, Ellipsis)
         one = seed * a + (1, Ellipsis)
         result = np.zeros(args.target_tensor.shape, args.target_tensor.dtype)