Enable mypy CI check (#926)

* Enable mypy CI check

- This enables the mypy CI check
- It also fixes the mypy issues introduced in the past few days.

Author: dstrain115

.github/workflows/ci.yaml

@@ -99,3 +99,20 @@ jobs:
           pip install --no-deps -e .
       - run: |
           check/pylint
+  
+  mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r dev_tools/requirements/envs/dev.env.txt
+          pip install --no-deps -e .
+      - run: |
+          check/mypy

dev_tools/conf/mypy.ini

@@ -3,6 +3,7 @@ show_error_codes = true
 plugins = duet.typing, numpy.typing.mypy_plugin
 allow_redefinition = true
 check_untyped_defs = true
+
 # Disabling function override checking
 # Qualtran has many places where kwargs are used
 # with the intention to override in subclasses in ways mypy does not like

qualtran/_infra/bloq.py

@@ -27,9 +27,11 @@
 
     from qualtran import (
         AddControlledT,
+        Adjoint,
         BloqBuilder,
         CompositeBloq,
         CtrlSpec,
+        GateWithRegisters,
         Register,
         Signature,
         Soquet,

qualtran/_infra/controlled.py

@@ -445,7 +445,7 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
         return self.ctrl_spec.wire_symbol(i, reg, idx)
 
     def adjoint(self) -> 'Bloq':
-        return self.subbloq.adjoint().controlled(self.ctrl_spec)
+        return self.subbloq.adjoint().controlled(ctrl_spec=self.ctrl_spec)
 
     def pretty_name(self) -> str:
         return f'C[{self.subbloq.pretty_name()}]'

qualtran/_infra/controlled_test.py

@@ -415,7 +415,7 @@ def test_controlled_diagrams():
 q: ──────X^0.25───''',
     )
 
-    ctrl_0_gate = XPowGate(0.25).controlled(CtrlSpec(cvs=0))
+    ctrl_0_gate = XPowGate(0.25).controlled(ctrl_spec=CtrlSpec(cvs=0))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(ctrl_0_gate.on_registers(**get_named_qubits(ctrl_0_gate.signature))),
         '''
@@ -424,7 +424,7 @@ def test_controlled_diagrams():
 q: ──────X^0.25───''',
     )
 
-    multi_ctrl_gate = XPowGate(0.25).controlled(CtrlSpec(cvs=[0, 1]))
+    multi_ctrl_gate = XPowGate(0.25).controlled(ctrl_spec=CtrlSpec(cvs=[0, 1]))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(multi_ctrl_gate.on_registers(**get_named_qubits(multi_ctrl_gate.signature))),
         '''
@@ -435,7 +435,7 @@ def test_controlled_diagrams():
 q: ─────────X^0.25───''',
     )
 
-    ctrl_bloq = Swap(2).controlled(CtrlSpec(cvs=[0, 1]))
+    ctrl_bloq = Swap(2).controlled(ctrl_spec=CtrlSpec(cvs=[0, 1]))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(ctrl_bloq.on_registers(**get_named_qubits(ctrl_bloq.signature))),
         '''

qualtran/_infra/gate_with_registers.py

@@ -15,6 +15,7 @@
 import abc
 from typing import (
     Any,
+    cast,
     Collection,
     Dict,
     Iterable,
@@ -357,8 +358,8 @@ def on_registers(
     ) -> cirq.Operation:
         return self.on(*merge_qubits(self.signature, **qubit_regs))
 
-    def __pow__(self, power: int) -> 'Bloq':
-        bloq = self if power > 0 else self.adjoint()
+    def __pow__(self, power: int) -> 'GateWithRegisters':
+        bloq = self if power > 0 else cast(GateWithRegisters, self.adjoint())
         if abs(power) == 1:
             return bloq
         if all(reg.side == Side.THRU for reg in self.signature):

qualtran/_infra/gate_with_registers_test.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Dict, TYPE_CHECKING
+from typing import Dict, Iterator, TYPE_CHECKING
 
 import cirq
 import numpy as np
@@ -46,7 +46,7 @@ def signature(self) -> Signature:
         regs = Signature([r1, r2, r3])
         return regs
 
-    def decompose_from_registers(self, *, context, **quregs) -> cirq.OP_TREE:
+    def decompose_from_registers(self, *, context, **quregs) -> Iterator[cirq.OP_TREE]:
         yield cirq.H.on_each(quregs['r1'])
         yield cirq.X.on_each(quregs['r2'])
         yield cirq.X.on_each(quregs['r3'])

qualtran/bloqs/arithmetic/addition.py

@@ -14,7 +14,19 @@
 import itertools
 import math
 from functools import cached_property
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import cirq
 import numpy as np
@@ -197,7 +209,7 @@ def _right_building_block(self, inp, out, anc, depth):
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         # reverse the order of qubits for big endian-ness.
         input_bits = quregs['a'][::-1]
         output_bits = quregs['b'][::-1]

qualtran/bloqs/arithmetic/comparison.py

@@ -90,7 +90,7 @@ def __pow__(self, power: int):
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         """Decomposes the gate into 4N And and And† operations for a T complexity of 4N.
 
         The decomposition proceeds from the most significant qubit -bit 0- to the least significant
@@ -217,7 +217,7 @@ def signature(self) -> Signature:
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         x, y, ancilla = quregs['x'], quregs['y'], quregs['ancilla']
         x_msb, x_lsb = x
         y_msb, y_lsb = y
@@ -310,7 +310,7 @@ def signature(self) -> Signature:
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         a = quregs['a']
         b = quregs['b']
         less_than = quregs['less_than']
@@ -358,7 +358,7 @@ def _sq_cmp() -> SingleQubitCompare:
 
 def _equality_with_zero(
     context: cirq.DecompositionContext, qubits: Sequence[cirq.Qid], z: cirq.Qid
-) -> cirq.OP_TREE:
+) -> Iterator[cirq.OP_TREE]:
     """Helper decomposition used in `LessThanEqual`"""
     if len(qubits) == 1:
         (q,) = qubits
@@ -451,7 +451,7 @@ def __pow__(self, power: int):
 
     def _decompose_via_tree(
         self, context: cirq.DecompositionContext, X: Sequence[cirq.Qid], Y: Sequence[cirq.Qid]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         if len(X) == 1:
             return
         if len(X) == 2:
@@ -467,7 +467,7 @@ def _decompose_via_tree(
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         lhs, rhs, (target,) = list(quregs['x']), list(quregs['y']), quregs['target']
 
         n = min(len(lhs), len(rhs))

qualtran/bloqs/arithmetic/hamming_weight.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import List, Set, TYPE_CHECKING
+from typing import Iterator, List, Set, TYPE_CHECKING
 
 import cirq
 from attrs import frozen
@@ -76,7 +76,7 @@ def _three_to_two_adder(self, a, b, c, out) -> cirq.OP_TREE:
 
     def _decompose_using_three_to_two_adders(
         self, x: List[cirq.Qid], junk: List[cirq.Qid], out: List[cirq.Qid]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         for out_idx in range(len(out)):
             y = []
             for in_idx in range(0, len(x) - 2, 2):
@@ -94,7 +94,7 @@ def _decompose_using_three_to_two_adders(
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         # Qubit order needs to be reversed because the registers store Big Endian representation
         # of integers.
         x: List[cirq.Qid] = [*quregs['x'][::-1]]