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

cirq/value/timestamp.py

@@ -11,8 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """A typed location in time that supports picosecond accuracy."""
 
+from __future__ import annotations
+
 from datetime import timedelta
 from typing import overload
 
@@ -45,23 +48,26 @@ class Timestamp:
         """The timestamp's location in picoseconds from arbitrary time zero."""
         return self._picos
 
-    def __add__(self, other) -> 'Timestamp':
+    def __add__(self, other) -> Timestamp:
         if isinstance(other, timedelta):
             return Timestamp(picos=self._picos + other.total_seconds() * 10**12)
         if not isinstance(other, Duration):
             return NotImplemented
         return Timestamp(picos=self._picos + other.total_picos())
 
-    def __radd__(self, other) -> 'Timestamp':
+    def __radd__(self, other) -> Timestamp:
         return self.__add__(other)
 
-    # pylint: disable=function-redefined
     @overload
-    def __sub__(self, other: 'Timestamp') -> Duration:
+    def __sub__(self, other: Timestamp) -> Duration:
+        pass
+
+    @overload
+    def __sub__(self, other: Duration) -> Timestamp:
         pass
 
     @overload
-    def __sub__(self, other: Duration) -> 'Timestamp':
+    def __sub__(self, other: timedelta) -> Timestamp:
         pass
 
     def __sub__(self, other):
@@ -73,8 +79,6 @@ class Timestamp:
             return Duration(picos=self._picos - other._picos)
         return NotImplemented
 
-    # pylint: enable=function-redefined
-
     def __eq__(self, other):
         if not isinstance(other, type(self)):
             return NotImplemented

cirq/value/timestamp_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 datetime import timedelta
 
 import pytest
@@ -20,7 +22,7 @@ import cirq
 from cirq import Duration, Timestamp
 
 
-def test_init():
+def test_init() -> None:
     assert Timestamp().raw_picos() == 0
     assert Timestamp(picos=513).raw_picos() == 513
     assert Timestamp(picos=-5).raw_picos() == -5
@@ -33,13 +35,13 @@ def test_init():
     assert isinstance(Timestamp(nanos=1.0).raw_picos(), float)
 
 
-def test_str():
+def test_str() -> None:
     assert str(Timestamp(picos=1000, nanos=1000)) == 't=1001000'
     assert str(Timestamp(nanos=5.0)) == 't=5000.0'
     assert str(Timestamp(picos=-100)) == 't=-100'
 
 
-def test_repr():
+def test_repr() -> None:
     a = Timestamp(picos=1000, nanos=1000)
     cirq.testing.assert_equivalent_repr(a)
     b = Timestamp(nanos=5.0)
@@ -48,14 +50,14 @@ def test_repr():
     cirq.testing.assert_equivalent_repr(c)
 
 
-def test_eq():
+def test_eq() -> None:
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(Timestamp(), Timestamp(picos=0), Timestamp(nanos=0.0))
     eq.add_equality_group(Timestamp(picos=1000), Timestamp(nanos=1))
     eq.make_equality_group(lambda: Timestamp(picos=-1))
 
 
-def test_cmp():
+def test_cmp() -> None:
     ordered_groups = [
         Timestamp(picos=-1),
         Timestamp(),
@@ -75,13 +77,13 @@ def test_cmp():
             assert (i >= j) == (a >= b)
             assert (i > j) == (a > b)
 
-    assert not (Timestamp() == 0)
+    assert not (Timestamp() == 0)  # noqa: SIM201
     assert Timestamp() != 0
-    assert not (Timestamp() == Duration())
+    assert not (Timestamp() == Duration())  # noqa: SIM201
     assert Timestamp() != Duration()
 
 
-def test_cmp_vs_other_type():
+def test_cmp_vs_other_type() -> None:
     with pytest.raises(TypeError):
         _ = Timestamp() < Duration()
     with pytest.raises(TypeError):
@@ -94,7 +96,7 @@ def test_cmp_vs_other_type():
         _ = Timestamp() > 0
 
 
-def test_add():
+def test_add() -> None:
     assert Timestamp(picos=1) + Duration(picos=2) == Timestamp(picos=3)
     assert Duration(picos=3) + Timestamp(picos=-5) == Timestamp(picos=-2)
 
@@ -109,7 +111,7 @@ def test_add():
         _ = Timestamp() + 1
 
 
-def test_sub():
+def test_sub() -> None:
     assert Timestamp() - Timestamp() == Duration()
     assert Timestamp(picos=1) - Timestamp(picos=2) == Duration(picos=-1)
     assert Timestamp(picos=5) - Duration(picos=2) == Timestamp(picos=3)
@@ -118,6 +120,6 @@ def test_sub():
     with pytest.raises(TypeError):
         _ = Duration() - Timestamp()
     with pytest.raises(TypeError):
-        _ = 1 - Timestamp()
+        _ = 1 - Timestamp()  # type: ignore[operator]
     with pytest.raises(TypeError):
-        _ = Timestamp() - 1
+        _ = Timestamp() - 1  # type: ignore[operator]

cirq/value/type_alias.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from __future__ import annotations
 
 import numpy as np
 import sympy
@@ -22,11 +22,11 @@ from cirq._doc import document
 """Supply aliases for commonly used types.
 """
 
-TParamKey = Union[str, sympy.Expr]
+TParamKey = str | sympy.Expr
 document(TParamKey, """A parameter that a parameter resolver may map to a value.""")
 
-TParamVal = Union[float, sympy.Expr]
+TParamVal = float | sympy.Expr
 document(TParamVal, """A value that a parameter resolver may return for a parameter.""")
 
-TParamValComplex = Union[complex, np.number, sympy.Expr]
+TParamValComplex = complex | np.number | sympy.Expr
 document(TParamValComplex, """A complex value that parameter resolvers may use for parameters.""")

cirq/value/value_equality_attr.py

@@ -11,9 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Defines `@cirq.value_equality`, for easy __eq__/__hash__ methods."""
 
-from typing import Any, Callable, Dict, Optional, overload, Union
+from __future__ import annotations
+
+from typing import Any, Callable, overload
 
 from typing_extensions import Protocol
 
@@ -83,7 +86,7 @@ def _value_equality_eq(self: _SupportsValueEquality, other: _SupportsValueEquali
 
 
 def _value_equality_ne(self: _SupportsValueEquality, other: _SupportsValueEquality) -> bool:
-    return not self == other
+    return not self == other  # noqa: SIM201
 
 
 def _value_equality_hash(self: _SupportsValueEquality) -> int:
@@ -110,7 +113,7 @@ def _value_equality_approx_eq(
     )
 
 
-def _value_equality_getstate(self: _SupportsValueEquality) -> Dict[str, Any]:
+def _value_equality_getstate(self: _SupportsValueEquality) -> dict[str, Any]:
     # clear cached hash value when pickling, see #6674
     state = self.__dict__
     hash_attr = _compat._method_cache_name(self.__hash__)
@@ -120,7 +123,6 @@ def _value_equality_getstate(self: _SupportsValueEquality) -> Dict[str, Any]:
     return state
 
 
-# pylint: disable=function-redefined
 @overload
 def value_equality(
     cls: type,
@@ -145,13 +147,13 @@ def value_equality(
 
 
 def value_equality(
-    cls: Optional[type] = None,
+    cls: type | None = None,
     *,
     unhashable: bool = False,
     distinct_child_types: bool = False,
     manual_cls: bool = False,
     approximate: bool = False,
-) -> Union[Callable[[type], type], type]:
+) -> Callable[[type], type] | type:
     """Implements __eq__/__ne__/__hash__ via a _value_equality_values_ method.
 
     _value_equality_values_ is a method that the decorated class must implement.
@@ -255,6 +257,3 @@ def value_equality(
         setattr(cls, '_approx_eq_', _value_equality_approx_eq)
 
     return cls
-
-
-# pylint: enable=function-redefined

cirq/value/value_equality_attr_test.py

@@ -11,6 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from __future__ import annotations
+
 import pytest
 
 import cirq
@@ -54,7 +57,7 @@ class BasicCb(BasicC):
     pass
 
 
-def test_value_equality_basic():
+def test_value_equality_basic() -> None:
 
     # Lookup works across equivalent types.
     v = {BasicC(1): 4, BasicCa(2): 5}
@@ -69,7 +72,7 @@ def test_value_equality_basic():
     eq.add_equality_group(BasicCa(3))
 
 
-def test_value_equality_manual():
+def test_value_equality_manual() -> None:
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(MasqueradePositiveD(3), BasicD(3))
     eq.add_equality_group(MasqueradePositiveD(4), MasqueradePositiveD(4), BasicD(4))
@@ -105,7 +108,7 @@ class UnhashableCb(UnhashableC):
     pass
 
 
-def test_value_equality_unhashable():
+def test_value_equality_unhashable() -> None:
     # Not possible to use as a dictionary key.
     with pytest.raises(TypeError, match='unhashable'):
         _ = {UnhashableC(1): 4}
@@ -143,7 +146,7 @@ class DistinctCb(DistinctC):
     pass
 
 
-def test_value_equality_distinct_child_types():
+def test_value_equality_distinct_child_types() -> None:
     # Lookup is distinct across child types.
     v = {DistinctC(1): 4, DistinctCa(1): 5, DistinctCb(1): 6}
     assert v[DistinctC(1)] == 4
@@ -168,7 +171,7 @@ class ApproxE:
         return self.x
 
 
-def test_value_equality_approximate():
+def test_value_equality_approximate() -> None:
     assert cirq.approx_eq(ApproxE(0.0), ApproxE(0.0), atol=0.1)
     assert cirq.approx_eq(ApproxE(0.0), ApproxE(0.2), atol=0.3)
     assert not cirq.approx_eq(ApproxE(0.0), ApproxE(0.2), atol=0.1)
@@ -187,13 +190,13 @@ class PeriodicF:
         return self.x % self.n
 
 
-def test_value_equality_approximate_specialized():
+def test_value_equality_approximate_specialized() -> None:
     assert PeriodicF(1, 4) != PeriodicF(5, 4)
     assert cirq.approx_eq(PeriodicF(1, 4), PeriodicF(5, 4), atol=0.1)
     assert not cirq.approx_eq(PeriodicF(1, 4), PeriodicF(6, 4), atol=0.1)
 
 
-def test_value_equality_approximate_not_supported():
+def test_value_equality_approximate_not_supported() -> None:
     assert not cirq.approx_eq(BasicC(0.0), BasicC(0.1), atol=0.2)
 
 
@@ -222,7 +225,7 @@ class ApproxGb(ApproxG):
     pass
 
 
-def test_value_equality_approximate_typing():
+def test_value_equality_approximate_typing() -> None:
     assert not cirq.approx_eq(ApproxE(0.0), PeriodicF(0.0, 1.0), atol=0.1)
     assert cirq.approx_eq(ApproxEa(0.0), ApproxEb(0.0), atol=0.1)
     assert cirq.approx_eq(ApproxG(0.0), ApproxG(0.0), atol=0.1)
@@ -230,7 +233,7 @@ def test_value_equality_approximate_typing():
     assert not cirq.approx_eq(ApproxG(0.0), ApproxGb(0.0), atol=0.1)
 
 
-def test_value_equality_forgot_method():
+def test_value_equality_forgot_method() -> None:
     with pytest.raises(TypeError, match='_value_equality_values_'):
 
         @cirq.value_equality
@@ -238,7 +241,7 @@ def test_value_equality_forgot_method():
             pass
 
 
-def test_bad_manual_cls_incompatible_args():
+def test_bad_manual_cls_incompatible_args() -> None:
     with pytest.raises(ValueError, match='incompatible'):
 
         @cirq.value_equality(manual_cls=True, distinct_child_types=True)
@@ -246,7 +249,7 @@ def test_bad_manual_cls_incompatible_args():
             pass
 
 
-def test_bad_manual_cls_forgot_method():
+def test_bad_manual_cls_forgot_method() -> None:
     with pytest.raises(TypeError, match='_value_equality_values_cls_'):
 
         @cirq.value_equality(manual_cls=True)

cirq/vis/density_matrix.py

@@ -14,7 +14,7 @@
 
 """Tool to visualize the magnitudes and phases in the density matrix"""
 
-from typing import Optional
+from __future__ import annotations
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -79,10 +79,10 @@ def _plot_element_of_density_matrix(ax, x, y, r, phase, show_rect=False, show_te
 
 def plot_density_matrix(
     matrix: np.ndarray,
-    ax: Optional[plt.Axes] = None,
+    ax: plt.Axes | None = None,
     *,
     show_text: bool = False,
-    title: Optional[str] = None,
+    title: str | None = None,
 ) -> plt.Axes:
     """Generates a plot for a given density matrix.
 

cirq/vis/density_matrix_test.py

@@ -11,8 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Tests for Density Matrix Plotter."""
 
+from __future__ import annotations
+
 import numpy as np
 import pytest
 from matplotlib import axis, lines, patches, pyplot as plt, spines, text
@@ -24,7 +27,7 @@ from cirq.vis.density_matrix import _plot_element_of_density_matrix, plot_densit
 @pytest.mark.usefixtures('closefigures')
 @pytest.mark.parametrize('show_text', [True, False])
 @pytest.mark.parametrize('size', [2, 4, 8, 16])
-def test_density_matrix_plotter(size, show_text):
+def test_density_matrix_plotter(size, show_text) -> None:
     matrix = cirq.testing.random_density_matrix(size)
     # Check that the title shows back up
     ax = plot_density_matrix(matrix, show_text=show_text, title='Test Density Matrix Plot')
@@ -48,16 +51,16 @@ def test_density_matrix_plotter(size, show_text):
 @pytest.mark.usefixtures('closefigures')
 @pytest.mark.parametrize('show_text', [True, False])
 @pytest.mark.parametrize('size', [2, 4, 8, 16])
-def test_density_matrix_circle_rectangle_sizes(size, show_text):
+def test_density_matrix_circle_rectangle_sizes(size, show_text) -> None:
     matrix = cirq.testing.random_density_matrix(size)
     ax = plot_density_matrix(matrix, show_text=show_text, title='Test Density Matrix Plot')
     # Check that the radius of all the circles in the matrix is correct
-    circles = [c for c in ax.get_children() if isinstance(c, patches.Circle)]
-    mean_radius = np.mean([c.radius for c in circles if c.fill])
+    circle_list = [c for c in ax.get_children() if isinstance(c, patches.Circle)]
+    mean_radius = np.mean([c.radius for c in circle_list if c.fill])
     mean_value = np.mean(np.abs(matrix))
-    circles = np.array(sorted(circles, key=lambda x: (x.fill, x.center[0], -x.center[1]))).reshape(
-        (2, size, size)
-    )
+    circles = np.array(
+        sorted(circle_list, key=lambda x: (x.fill, x.center[0], -x.center[1]))
+    ).reshape((2, size, size))
     for i in range(size):
         for j in range(size):
             assert np.isclose(
@@ -65,15 +68,15 @@ def test_density_matrix_circle_rectangle_sizes(size, show_text):
             )
 
     # Check that all the rectangles are of the right height, and only on the diagonal elements
-    rects = [
+    rect_list = [
         r
         for r in ax.get_children()
         if isinstance(r, patches.Rectangle) and r.get_alpha() is not None
     ]
-    assert len(rects) == size
-    mean_size = np.mean([r.get_height() for r in rects])
+    assert len(rect_list) == size
+    mean_size = np.mean([r.get_height() for r in rect_list])
     mean_value = np.trace(np.abs(matrix)) / size
-    rects = np.array(sorted(rects, key=lambda x: x.get_x()))
+    rects = np.array(sorted(rect_list, key=lambda x: x.get_x()))
     for i in range(size):
         # Ensuring that the rectangle is the right height
         assert np.isclose(np.abs(matrix[i, i]) * mean_size / mean_value, rects[i].get_height())
@@ -90,7 +93,7 @@ def test_density_matrix_circle_rectangle_sizes(size, show_text):
 @pytest.mark.usefixtures('closefigures')
 @pytest.mark.parametrize('show_text', [True, False])
 @pytest.mark.parametrize('size', [2, 4, 8, 16])
-def test_density_matrix_sizes_upper_bounds(size, show_text):
+def test_density_matrix_sizes_upper_bounds(size, show_text) -> None:
     matrix = cirq.testing.random_density_matrix(size)
     ax = plot_density_matrix(matrix, show_text=show_text, title='Test Density Matrix Plot')
 
@@ -113,7 +116,7 @@ def test_density_matrix_sizes_upper_bounds(size, show_text):
 @pytest.mark.usefixtures('closefigures')
 @pytest.mark.parametrize('show_rect', [True, False])
 @pytest.mark.parametrize('value', [0.0, 1.0, 0.5 + 0.3j, 0.2 + 0.1j, 0.5 + 0.5j])
-def test_density_element_plot(value, show_rect):
+def test_density_element_plot(value, show_rect) -> None:
     _, ax = plt.subplots(figsize=(10, 10))
     _plot_element_of_density_matrix(
         ax, 0, 0, np.abs(value), np.angle(value), show_rect=False, show_text=False
@@ -122,9 +125,9 @@ def test_density_element_plot(value, show_rect):
     plotted_lines = [c for c in ax.get_children() if isinstance(c, lines.Line2D)]
     assert len(plotted_lines) == 1
     line_position = plotted_lines[0].get_xydata()
-    angle = np.arctan(
-        (line_position[1, 1] - line_position[0, 1]) / (line_position[1, 0] - line_position[0, 0])
-    )
+    numerator = line_position[1, 1] - line_position[0, 1]  # type: ignore
+    denumerator = line_position[1, 0] - line_position[0, 0]  # type: ignore
+    angle = np.arctan(numerator / denumerator)
     assert np.isclose(np.angle(value), angle)
     # Check if the circles are the right size ratio, given the value of the element
     circles_in = [c for c in ax.get_children() if isinstance(c, patches.Circle) and c.fill]
@@ -151,6 +154,6 @@ def test_density_element_plot(value, show_rect):
         np.random.random((4, 8)) * np.exp(np.random.random((4, 8)) * 2 * np.pi * 1j),
     ],
 )
-def test_density_matrix_type_error(matrix):
+def test_density_matrix_type_error(matrix) -> None:
     with pytest.raises(ValueError, match="Incorrect shape for density matrix:*"):
         plot_density_matrix(matrix)

cirq/vis/heatmap.py

@@ -16,19 +16,7 @@ from __future__ import annotations
 
 import copy
 from dataclasses import astuple, dataclass
-from typing import (
-    Any,
-    cast,
-    Dict,
-    List,
-    Mapping,
-    Optional,
-    overload,
-    Sequence,
-    SupportsFloat,
-    Tuple,
-    Union,
-)
+from typing import Any, cast, Mapping, overload, Sequence, SupportsFloat
 
 import matplotlib as mpl
 import matplotlib.collections as mpl_collections
@@ -39,9 +27,9 @@ from mpl_toolkits import axes_grid1
 from cirq.devices import grid_qubit
 from cirq.vis import vis_utils
 
-QubitTuple = Tuple[grid_qubit.GridQubit, ...]
+QubitTuple = tuple[grid_qubit.GridQubit, ...]
 
-Polygon = Sequence[Tuple[float, float]]
+Polygon = Sequence[tuple[float, float]]
 
 
 @dataclass
@@ -71,13 +59,12 @@ class PolygonUnit:
     polygon: Polygon
     value: float
     center: Point
-    annot: Optional[str]
+    annot: str | None
 
 
 class Heatmap:
     """Distribution of a value in 2D qubit lattice as a color map."""
 
-    # pylint: disable=function-redefined
     @overload
     def __init__(self, value_map: Mapping[QubitTuple, SupportsFloat], **kwargs):
         pass
@@ -88,9 +75,9 @@ class Heatmap:
 
     def __init__(
         self,
-        value_map: Union[
-            Mapping[QubitTuple, SupportsFloat], Mapping[grid_qubit.GridQubit, SupportsFloat]
-        ],
+        value_map: (
+            Mapping[QubitTuple, SupportsFloat] | Mapping[grid_qubit.GridQubit, SupportsFloat]
+        ),
         **kwargs,
     ):
         """2D qubit grid Heatmaps
@@ -131,7 +118,7 @@ class Heatmap:
         }
         self._validate_kwargs(kwargs)
         if '_config' not in self.__dict__:
-            self._config: Dict[str, Any] = {}
+            self._config: dict[str, Any] = {}
         self._config.update(
             {
                 "plot_colorbar": True,
@@ -144,7 +131,7 @@ class Heatmap:
         )
         self._config.update(kwargs)
 
-    def _extra_valid_kwargs(self) -> List[str]:
+    def _extra_valid_kwargs(self) -> list[str]:
         return []
 
     def _validate_kwargs(self, kwargs) -> None:
@@ -173,13 +160,13 @@ class Heatmap:
             invalid_args = ", ".join([k for k in kwargs if k not in valid_kwargs])
             raise ValueError(f"Received invalid argument(s): {invalid_args}")
 
-    def update_config(self, **kwargs) -> 'Heatmap':
+    def update_config(self, **kwargs) -> Heatmap:
         """Add/Modify **kwargs args passed during initialisation."""
         self._validate_kwargs(kwargs)
         self._config.update(kwargs)
         return self
 
-    def _qubits_to_polygon(self, qubits: QubitTuple) -> Tuple[Polygon, Point]:
+    def _qubits_to_polygon(self, qubits: QubitTuple) -> tuple[Polygon, Point]:
         qubit = qubits[0]
         x, y = float(qubit.row), float(qubit.col)
         return (
@@ -187,7 +174,7 @@ class Heatmap:
             Point(y, x),
         )
 
-    def _get_annotation_value(self, key, value) -> Optional[str]:
+    def _get_annotation_value(self, key, value) -> str | None:
         if self._config.get('annotation_map'):
             return self._config['annotation_map'].get(key)
         elif self._config.get('annotation_format'):
@@ -198,8 +185,8 @@ class Heatmap:
         else:
             return None
 
-    def _get_polygon_units(self) -> List[PolygonUnit]:
-        polygon_unit_list: List[PolygonUnit] = []
+    def _get_polygon_units(self) -> list[PolygonUnit]:
+        polygon_unit_list: list[PolygonUnit] = []
         for qubits, value in sorted(self._value_map.items()):
             polygon, center = self._qubits_to_polygon(qubits)
             polygon_unit_list.append(
@@ -231,7 +218,7 @@ class Heatmap:
 
     def _write_annotations(
         self,
-        centers_and_annot: List[Tuple[Point, Optional[str]]],
+        centers_and_annot: list[tuple[Point, str | None]],
         collection: mpl_collections.Collection,
         ax: plt.Axes,
     ) -> None:
@@ -245,7 +232,7 @@ class Heatmap:
             x, y = center
             face_luminance = vis_utils.relative_luminance(facecolor)
             text_color = 'black' if face_luminance > 0.4 else 'white'
-            text_kwargs: Dict[str, Any] = dict(color=text_color, ha="center", va="center")
+            text_kwargs: dict[str, Any] = dict(color=text_color, ha="center", va="center")
             text_kwargs.update(self._config.get('annotation_text_kwargs', {}))
             ax.text(x, y, annotation, **text_kwargs)
 
@@ -286,8 +273,8 @@ class Heatmap:
         return collection
 
     def plot(
-        self, ax: Optional[plt.Axes] = None, **kwargs: Any
-    ) -> Tuple[plt.Axes, mpl_collections.Collection]:
+        self, ax: plt.Axes | None = None, **kwargs: Any
+    ) -> tuple[plt.Axes, mpl_collections.Collection]:
         """Plots the heatmap on the given Axes.
         Args:
             ax: the Axes to plot on. If not given, a new figure is created,
@@ -358,13 +345,13 @@ class TwoQubitInteractionHeatmap(Heatmap):
                 coupler_margin: float, default = 0.03
                 coupler_width: float, default = 0.6
         """
-        self._config: Dict[str, Any] = {"coupler_margin": 0.03, "coupler_width": 0.6}
+        self._config: dict[str, Any] = {"coupler_margin": 0.03, "coupler_width": 0.6}
         super().__init__(value_map, **kwargs)
 
-    def _extra_valid_kwargs(self) -> List[str]:
+    def _extra_valid_kwargs(self) -> list[str]:
         return ["coupler_margin", "coupler_width"]
 
-    def _qubits_to_polygon(self, qubits: QubitTuple) -> Tuple[Polygon, Point]:
+    def _qubits_to_polygon(self, qubits: QubitTuple) -> tuple[Polygon, Point]:
         coupler_margin = self._config["coupler_margin"]
         coupler_width = self._config["coupler_width"]
         cwidth = coupler_width / 2.0
@@ -403,8 +390,8 @@ class TwoQubitInteractionHeatmap(Heatmap):
         return (polygon, Point((col1 + col2) / 2.0, (row1 + row2) / 2.0))
 
     def plot(
-        self, ax: Optional[plt.Axes] = None, **kwargs: Any
-    ) -> Tuple[plt.Axes, mpl_collections.Collection]:
+        self, ax: plt.Axes | None = None, **kwargs: Any
+    ) -> tuple[plt.Axes, mpl_collections.Collection]:
         """Plots the heatmap on the given Axes.
         Args:
             ax: the Axes to plot on. If not given, a new figure is created,
@@ -422,7 +409,7 @@ class TwoQubitInteractionHeatmap(Heatmap):
         original_config = copy.deepcopy(self._config)
         self.update_config(**kwargs)
         qubits = set([q for qubits in self._value_map.keys() for q in qubits])
-        collection_options: Dict[str, Any] = {"cmap": "binary"}
+        collection_options: dict[str, Any] = {"cmap": "binary"}
         highlighted_qubits = frozenset(kwargs.get("highlighted_qubits", ()))
         if not highlighted_qubits:
             collection_options.update(

cirq/vis/heatmap_test.py

@@ -11,8 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Tests for Heatmap."""
 
+from __future__ import annotations
+
 import pathlib
 import shutil
 import string