cirq-core 1.6.0.dev20250516060316__py3-none-any.whl → 1.6.0.dev20250517012327__py3-none-any.whl

cirq/__init__.py

@@ -397,6 +397,7 @@ from cirq.transformers import (
     single_qubit_matrix_to_phxz as single_qubit_matrix_to_phxz,
     single_qubit_op_to_framed_phase_form as single_qubit_op_to_framed_phase_form,
     stratified_circuit as stratified_circuit,
+    symbolize_single_qubit_gates_by_indexed_tags as symbolize_single_qubit_gates_by_indexed_tags,
     synchronize_terminal_measurements as synchronize_terminal_measurements,
     TRANSFORMER as TRANSFORMER,
     TransformerContext as TransformerContext,

cirq/_import_test.py

@@ -17,9 +17,9 @@ from __future__ import annotations
 from cirq import _import
 
 
-def test_lazy_loader():
+def test_lazy_loader() -> None:
     linalg = _import.LazyLoader("linalg", globals(), "scipy.linalg")
-    linalg.fun = 1
+    linalg.fun = 1  # type: ignore[attr-defined]
     assert linalg._module is None
     assert "linalg" not in linalg.__dict__
 

cirq/_version.py

@@ -28,4 +28,4 @@ if sys.version_info < (3, 11, 0):  # pragma: no cover
         'of cirq (e.g. "python -m pip install cirq==1.5.0")'
     )
 
-__version__ = "1.6.0.dev20250516060316"
+__version__ = "1.6.0.dev20250517012327"

cirq/_version_test.py

@@ -3,4 +3,4 @@ import cirq
 
 
 def test_version() -> None:
-    assert cirq.__version__ == "1.6.0.dev20250516060316"
+    assert cirq.__version__ == "1.6.0.dev20250517012327"

cirq/interop/quirk/cells/parse_test.py

@@ -19,7 +19,7 @@ import sympy
 from cirq.interop.quirk.cells.parse import parse_complex, parse_formula, parse_matrix
 
 
-def test_parse_matrix():
+def test_parse_matrix() -> None:
     s = np.sqrt(0.5)
     np.testing.assert_allclose(
         parse_matrix('{{√½,√½},{-√½,√½}}'), np.array([[s, s], [-s, s]]), atol=1e-8
@@ -32,7 +32,7 @@ def test_parse_matrix():
     )
 
 
-def test_parse_matrix_failures():
+def test_parse_matrix_failures() -> None:
     with pytest.raises(ValueError, match='Not surrounded by {{}}'):
         _ = parse_matrix('1')
     with pytest.raises(ValueError, match='Not surrounded by {{}}'):
@@ -45,7 +45,7 @@ def test_parse_matrix_failures():
         _ = parse_matrix('{{x}}')
 
 
-def test_parse_real_formula():
+def test_parse_real_formula() -> None:
     t = sympy.Symbol('t')
     assert parse_formula('1/2') == 0.5
     assert parse_formula('t*t + ln(t)') == t * t + sympy.ln(t)
@@ -58,12 +58,12 @@ def test_parse_real_formula():
         _ = parse_formula('i')
 
 
-def test_parse_formula_failures():
+def test_parse_formula_failures() -> None:
     with pytest.raises(TypeError, match='formula must be a string'):
-        _ = parse_formula(2)
+        _ = parse_formula(2)  # type: ignore[arg-type]
 
     with pytest.raises(TypeError, match='formula must be a string'):
-        _ = parse_formula([])
+        _ = parse_formula([])  # type: ignore[arg-type]
 
     with pytest.raises(ValueError, match='Unrecognized token'):
         _ = parse_formula('5*__**DSA **)SADD')
@@ -72,7 +72,7 @@ def test_parse_formula_failures():
         _ = parse_formula('5*x')
 
 
-def test_parse_complex():
+def test_parse_complex() -> None:
     assert parse_complex('0') == 0
     assert parse_complex('1') == 1
     assert parse_complex('i') == 1j
@@ -84,7 +84,7 @@ def test_parse_complex():
     np.testing.assert_allclose(parse_complex('exp 2'), np.e**2, atol=1e-8)
 
 
-def test_parse_complex_raw_cases_from_quirk():
+def test_parse_complex_raw_cases_from_quirk() -> None:
     assert parse_complex("0") == 0
     assert parse_complex("1") == 1
     assert parse_complex("-1") == -1
@@ -116,7 +116,7 @@ def test_parse_complex_raw_cases_from_quirk():
     np.testing.assert_allclose(parse_complex("2       pi"), 2 * np.pi)
 
 
-def test_parse_complex_expression_cases_from_quirk():
+def test_parse_complex_expression_cases_from_quirk() -> None:
     np.testing.assert_allclose(parse_complex("1/3"), 1 / 3)
     np.testing.assert_allclose(parse_complex("2/3/5"), (2 / 3) / 5)
     np.testing.assert_allclose(parse_complex("2/3/5*7/13"), ((((2 / 3) / 5)) * 7) / 13)
@@ -152,7 +152,7 @@ def test_parse_complex_expression_cases_from_quirk():
     np.testing.assert_allclose(parse_complex("cos(acos(0.5))"), 0.5, atol=1e-8)
 
 
-def test_parse_complex_expression_failures():
+def test_parse_complex_expression_failures() -> None:
     with pytest.raises(ValueError, match='Incomplete expression'):
         _ = parse_formula('(')
     with pytest.raises(ValueError, match=r"unmatched '\)'"):

cirq/testing/equals_tester_test.py

@@ -19,7 +19,7 @@ import pytest
 from cirq.testing.equals_tester import EqualsTester
 
 
-def test_add_equality_group_correct():
+def test_add_equality_group_correct() -> None:
     eq = EqualsTester()
 
     eq.add_equality_group(fractions.Fraction(1, 1))
@@ -39,7 +39,7 @@ def test_add_equality_group_correct():
     eq.add_equality_group('unrelated')
 
 
-def test_assert_make_equality_group():
+def test_assert_make_equality_group() -> None:
     eq = EqualsTester()
 
     with pytest.raises(AssertionError, match="can't be in the same"):
@@ -57,20 +57,20 @@ def test_assert_make_equality_group():
         eq.make_equality_group(lambda: 3)
 
 
-def test_add_equality_group_not_equivalent():
+def test_add_equality_group_not_equivalent() -> None:
     eq = EqualsTester()
     with pytest.raises(AssertionError, match="can't be in the same"):
         eq.add_equality_group(1, 2)
 
 
-def test_add_equality_group_not_disjoint():
+def test_add_equality_group_not_disjoint() -> None:
     eq = EqualsTester()
     eq.add_equality_group(1)
     with pytest.raises(AssertionError, match="can't be in different"):
         eq.add_equality_group(1)
 
 
-def test_add_equality_group_bad_hash():
+def test_add_equality_group_bad_hash() -> None:
     class KeyHash:
         def __init__(self, k, h):
             self._k = k
@@ -94,7 +94,7 @@ def test_add_equality_group_bad_hash():
         eq.add_equality_group(KeyHash('c', 2), KeyHash('c', 3))
 
 
-def test_add_equality_group_exception_hash():
+def test_add_equality_group_exception_hash() -> None:
     class FailHash:
         def __hash__(self):
             raise ValueError('injected failure')
@@ -104,7 +104,7 @@ def test_add_equality_group_exception_hash():
         eq.add_equality_group(FailHash())
 
 
-def test_fails_when_forgot_type_check():
+def test_fails_when_forgot_type_check() -> None:
     eq = EqualsTester()
 
     class NoTypeCheckEqualImplementation:
@@ -124,11 +124,12 @@ def test_fails_when_forgot_type_check():
         eq.add_equality_group(NoTypeCheckEqualImplementation())
 
 
-def test_fails_when_equal_to_everything():
+def test_fails_when_equal_to_everything() -> None:
     eq = EqualsTester()
 
     class AllEqual:
-        __hash__ = None
+        def __hash__(self) -> int:
+            return 0
 
         def __eq__(self, other):
             return True
@@ -136,11 +137,12 @@ def test_fails_when_equal_to_everything():
         def __ne__(self, other):
             return False
 
+    assert hash(AllEqual()) == 0
     with pytest.raises(AssertionError, match="can't be in different"):
         eq.add_equality_group(AllEqual())
 
 
-def test_fails_hash_is_default_and_inconsistent():
+def test_fails_hash_is_default_and_inconsistent() -> None:
     eq = EqualsTester()
 
     class DefaultHashImplementation:
@@ -161,7 +163,7 @@ def test_fails_hash_is_default_and_inconsistent():
         eq.make_equality_group(DefaultHashImplementation)
 
 
-def test_fails_when_ne_is_inconsistent():
+def test_fails_when_ne_is_inconsistent() -> None:
     eq = EqualsTester()
 
     class InconsistentNeImplementation:
@@ -185,7 +187,7 @@ def test_fails_when_ne_is_inconsistent():
         eq.make_equality_group(InconsistentNeImplementation)
 
 
-def test_fails_when_ne_is_inconsistent_due_to_not_implemented():
+def test_fails_when_ne_is_inconsistent_due_to_not_implemented() -> None:
     eq = EqualsTester()
 
     class InconsistentNeImplementation:
@@ -207,7 +209,7 @@ def test_fails_when_ne_is_inconsistent_due_to_not_implemented():
         eq.make_equality_group(InconsistentNeImplementation)
 
 
-def test_fails_when_not_reflexive():
+def test_fails_when_not_reflexive() -> None:
     eq = EqualsTester()
 
     class NotReflexiveImplementation:
@@ -226,7 +228,7 @@ def test_fails_when_not_reflexive():
         eq.add_equality_group(NotReflexiveImplementation())
 
 
-def test_fails_when_not_commutative():
+def test_fails_when_not_commutative() -> None:
     eq = EqualsTester()
 
     class NotCommutativeImplementation:
@@ -248,14 +250,14 @@ def test_fails_when_not_commutative():
         eq.add_equality_group(NotCommutativeImplementation(1), NotCommutativeImplementation(0))
 
 
-def test_works_on_types():
+def test_works_on_types() -> None:
     eq = EqualsTester()
     eq.add_equality_group(object)
     eq.add_equality_group(int)
     eq.add_equality_group(object())
 
 
-def test_returns_not_implemented_for_other_types():
+def test_returns_not_implemented_for_other_types() -> None:
     # First we demonstrate an example of the problem.
 
     # FirstClass is the class that is broken.
@@ -323,7 +325,7 @@ def test_returns_not_implemented_for_other_types():
     eq.add_equality_group(ThirdClass("a"), ThirdClass("a"))
 
 
-def test_not_implemented_error():
+def test_not_implemented_error() -> None:
     # Common bug is to return NotImplementedError instead of NotImplemented.
     class NotImplementedErrorCase:
         def __init__(self, val):

cirq/transformers/__init__.py

@@ -135,6 +135,11 @@ from cirq.transformers.transformer_primitives import (
     unroll_circuit_op_greedy_frontier as unroll_circuit_op_greedy_frontier,
 )
 
+from cirq.transformers.symbolize import (
+    SymbolizeTag as SymbolizeTag,
+    symbolize_single_qubit_gates_by_indexed_tags as symbolize_single_qubit_gates_by_indexed_tags,
+)
+
 from cirq.transformers.gauge_compiling import (
     CZGaugeTransformer as CZGaugeTransformer,
     ConstantGauge as ConstantGauge,

cirq/transformers/symbolize.py

@@ -0,0 +1,101 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Transformers that symbolize operations."""
+
+import re
+from typing import Hashable, Optional, TYPE_CHECKING
+
+import attrs
+import sympy
+from attrs import validators
+
+from cirq import ops
+from cirq.transformers import transformer_api, transformer_primitives
+
+if TYPE_CHECKING:
+    import cirq
+
+
+@attrs.frozen
+class SymbolizeTag:
+    prefix: str = attrs.field(
+        validator=validators.and_(validators.instance_of(str), validators.min_len(1))
+    )
+
+
+@transformer_api.transformer
+def symbolize_single_qubit_gates_by_indexed_tags(
+    circuit: 'cirq.AbstractCircuit',
+    *,
+    context: Optional['cirq.TransformerContext'] = None,
+    symbolize_tag: SymbolizeTag = SymbolizeTag(prefix="TO-PHXZ"),
+) -> 'cirq.Circuit':
+    """Symbolizes single qubit operations by indexed tags prefixed by symbolize_tag.prefix.
+
+    Example:
+        >>> from cirq import transformers
+        >>> q0, q1 = cirq.LineQubit.range(2)
+        >>> c = cirq.Circuit(
+        ...         cirq.X(q0).with_tags("phxz_0"),
+        ...         cirq.CZ(q0,q1),
+        ...         cirq.Y(q0).with_tags("phxz_1"),
+        ...         cirq.X(q0))
+        >>> print(c)
+        0: ───X[phxz_0]───@───Y[phxz_1]───X───
+                          │
+        1: ───────────────@───────────────────
+        >>> new_circuit = cirq.symbolize_single_qubit_gates_by_indexed_tags(
+        ...     c, symbolize_tag=transformers.SymbolizeTag(prefix="phxz"))
+        >>> print(new_circuit)
+        0: ───PhXZ(a=a0,x=x0,z=z0)───@───PhXZ(a=a1,x=x1,z=z1)───X───
+                                     │
+        1: ──────────────────────────@──────────────────────────────
+
+    Args:
+        circuit: Input circuit to apply the transformations on. The input circuit is not mutated.
+        context: `cirq.TransformerContext` storing common configurable options for transformers.
+        symbolize_tag: The tag info used to symbolize the phxz gate. Prefix is required.
+
+    Returns:
+        Copy of the transformed input circuit.
+    """
+
+    def _map_func(op: 'cirq.Operation', _):
+        """Maps an op with tag `{tag_prefix}_i` to a symbolzied `PhasedXZGate(xi,zi,ai)`."""
+        tags: set[Hashable] = set(op.tags)
+        tag_id: None | int = None
+        for tag in tags:
+            if re.fullmatch(f"{symbolize_tag.prefix}_\\d+", str(tag)):
+                if tag_id is None:
+                    tag_id = int(str(tag).rsplit("_", maxsplit=-1)[-1])
+                else:
+                    raise ValueError(f"Multiple tags are prefixed with {symbolize_tag.prefix}.")
+        if tag_id is None:
+            return op
+        tags.remove(f"{symbolize_tag.prefix}_{tag_id}")
+        phxz_params = {
+            "x_exponent": sympy.Symbol(f"x{tag_id}"),
+            "z_exponent": sympy.Symbol(f"z{tag_id}"),
+            "axis_phase_exponent": sympy.Symbol(f"a{tag_id}"),
+        }
+
+        return ops.PhasedXZGate(**phxz_params).on(*op.qubits).with_tags(*tags)
+
+    return transformer_primitives.map_operations(
+        circuit.freeze(),
+        _map_func,
+        deep=context.deep if context else False,
+        tags_to_ignore=context.tags_to_ignore if context else [],
+    ).unfreeze(copy=False)

cirq/transformers/symbolize_test.py

@@ -0,0 +1,60 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import pytest
+import sympy
+
+import cirq
+from cirq.transformers.symbolize import SymbolizeTag
+
+
+def test_symbolize_single_qubit_gates_by_indexed_tags_success():
+    q = cirq.NamedQubit("a")
+    input_circuit = cirq.Circuit(
+        cirq.X(q).with_tags("phxz_1"), cirq.Y(q).with_tags("tag1"), cirq.Z(q).with_tags("phxz_0")
+    )
+    output_circuit = cirq.symbolize_single_qubit_gates_by_indexed_tags(
+        input_circuit, symbolize_tag=SymbolizeTag(prefix="phxz")
+    )
+    cirq.testing.assert_same_circuits(
+        output_circuit,
+        cirq.Circuit(
+            cirq.PhasedXZGate(
+                x_exponent=sympy.Symbol("x1"),
+                z_exponent=sympy.Symbol("z1"),
+                axis_phase_exponent=sympy.Symbol("a1"),
+            ).on(q),
+            cirq.Y(q).with_tags("tag1"),
+            cirq.PhasedXZGate(
+                x_exponent=sympy.Symbol("x0"),
+                z_exponent=sympy.Symbol("z0"),
+                axis_phase_exponent=sympy.Symbol("a0"),
+            ).on(q),
+        ),
+    )
+
+
+def test_symbolize_single_qubit_gates_by_indexed_tags_multiple_tags():
+    q = cirq.NamedQubit("a")
+    input_circuit = cirq.Circuit(cirq.X(q).with_tags("TO-PHXZ_0", "TO-PHXZ_2"))
+
+    with pytest.raises(ValueError, match="Multiple tags are prefixed with TO-PHXZ."):
+        cirq.symbolize_single_qubit_gates_by_indexed_tags(input_circuit)
+
+
+def test_symbolize_tag_invalid_prefix():
+    with pytest.raises(ValueError, match="Length of 'prefix' must be >= 1: 0"):
+        SymbolizeTag(prefix="")
+    with pytest.raises(TypeError, match="'prefix' must be <class 'str'>"):
+        SymbolizeTag(prefix=[1])

cirq/value/random_state_test.py

@@ -17,7 +17,7 @@ import numpy as np
 import cirq
 
 
-def test_parse_random_state():
+def test_parse_random_state() -> None:
     global_state = np.random.get_state()
 
     def rand(prng):
@@ -34,11 +34,11 @@ def test_parse_random_state():
     eq.add_equality_group(*vals)
 
     seed = np.random.randint(2**31)
-    prngs = [
+    prngs1 = [
         np.random.RandomState(seed),
         cirq.value.parse_random_state(np.random.RandomState(seed)),
         cirq.value.parse_random_state(seed),
     ]
-    vals = [prng.rand() for prng in prngs]
+    vals = [prng.rand() for prng in prngs1]
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(*vals)

cirq/value/timestamp.py

@@ -64,6 +64,10 @@ class Timestamp:
     def __sub__(self, other: Duration) -> 'Timestamp':
         pass
 
+    @overload
+    def __sub__(self, other: timedelta) -> 'Timestamp':
+        pass
+
     def __sub__(self, other):
         if isinstance(other, Duration):
             return Timestamp(picos=self._picos - other.total_picos())

cirq/value/timestamp_test.py

@@ -20,7 +20,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 +33,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 +48,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(),
@@ -81,7 +81,7 @@ def test_cmp():
     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 +94,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 +109,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 +118,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/vis/density_matrix_test.py

@@ -24,7 +24,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 +48,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 +65,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 +90,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 +113,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 +122,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 +151,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)