Mypy issues part 5 (#915)

* Mypy issues part 5

- Solves almost all of the remaining mypy issues
- Mostly confusion about symbolic expressions and whether
things should be numpy arrays or not

Author: dstrain115

qualtran/_infra/adjoint.py

@@ -199,7 +199,7 @@ def _t_complexity_(self):
             return NotImplemented
 
         try:
-            return self.subbloq._t_complexity_(adjoint=True)
+            return self.subbloq._t_complexity_(adjoint=True)  # type: ignore[call-arg]
         except TypeError as e:
             if 'adjoint' in str(e):
                 return self.subbloq._t_complexity_()

qualtran/_infra/bloq.py

@@ -180,7 +180,9 @@ def adjoint(self) -> 'Bloq':
 
         return Adjoint(self)
 
-    def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
+    def on_classical_vals(
+        self, **vals: Union['sympy.Symbol', 'ClassicalValT']
+    ) -> Dict[str, 'ClassicalValT']:
         """How this bloq operates on classical data.
 
         Override this method if your bloq represents classical, reversible logic. For example:

qualtran/_infra/bloq_example.py

@@ -90,8 +90,11 @@ def bloq_example(
 
 
 def bloq_example(
-    _func: Optional[Callable[[], _BloqType]] = None, *, generalizer: _GeneralizerType = lambda x: x
-) -> BloqExample[_BloqType]:
+    _func: Optional[Callable[[], _BloqType]] = None,
+    *,
+    generalizer: _GeneralizerType = lambda x: x,
+    **kwargs: Any,
+) -> Union[Callable[[Callable[[], _BloqType]], BloqExample[_BloqType]], BloqExample[_BloqType]]:
     """Decorator to turn a function into a `BloqExample`.
 
     This will set `name` to the name of the function and `bloq_cls` according to the return-type

qualtran/_infra/composite_bloq.py

@@ -16,12 +16,14 @@
 from functools import cached_property
 from typing import (
     Callable,
+    cast,
     Dict,
     FrozenSet,
     Hashable,
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     overload,
     Sequence,
@@ -176,7 +178,9 @@ def from_cirq_circuit(cls, circuit: 'cirq.Circuit') -> 'CompositeBloq':
 
         return cirq_optree_to_cbloq(circuit)
 
-    def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT']:
+    def on_classical_vals(
+        self, **vals: Union[sympy.Symbol, 'ClassicalValT']
+    ) -> Dict[str, 'ClassicalValT']:
         """Support classical data by recursing into the composite bloq."""
         from qualtran.simulation.classical_sim import call_cbloq_classically
 
@@ -606,7 +610,7 @@ def _reg_to_soq(
 
 def _process_soquets(
     registers: Iterable[Register],
-    in_soqs: Dict[str, SoquetT],
+    in_soqs: Mapping[str, SoquetInT],
     debug_str: str,
     func: Callable[[Soquet, Register, Tuple[int, ...]], None],
 ) -> None:
@@ -631,6 +635,7 @@ def _process_soquets(
             the incoming, indexed soquet as well as the register and (left-)index it
             has been mapped to.
     """
+    unchecked_names: Set[str] = set(in_soqs.keys())
     for reg in registers:
         try:
             # if we want fancy indexing (which we do), we need numpy
@@ -639,7 +644,7 @@ def _process_soquets(
         except KeyError:
             raise BloqError(f"{debug_str} requires a Soquet named `{reg.name}`.") from None
 
-        del in_soqs[reg.name]  # so we can check for surplus arguments.
+        unchecked_names.remove(reg.name)  # so we can check for surplus arguments.
 
         for li in reg.all_idxs():
             idxed_soq = in_soq[li]
@@ -652,9 +657,8 @@ def _process_soquets(
                 raise BloqError(
                     f"{debug_str} register dtypes are not consistent {extra_str}."
                 ) from None
-
-    if in_soqs:
-        raise BloqError(f"{debug_str} does not accept Soquets: {in_soqs.keys()}.") from None
+    if unchecked_names:
+        raise BloqError(f"{debug_str} does not accept Soquets: {unchecked_names}.") from None
 
 
 def _map_soqs(
@@ -965,7 +969,7 @@ def add(self, bloq: Bloq, **in_soqs: SoquetInT):
         return outs
 
     def _add_binst(
-        self, binst: BloqInstance, in_soqs: Dict[str, SoquetInT]
+        self, binst: BloqInstance, in_soqs: Mapping[str, SoquetInT]
     ) -> Iterator[Tuple[str, SoquetT]]:
         """Add a bloq instance.
 
@@ -1010,7 +1014,8 @@ def add_from(self, bloq: Bloq, **in_soqs: SoquetInT) -> Tuple[SoquetT, ...]:
 
         # Initial mapping of LeftDangle according to user-provided in_soqs.
         soq_map: List[Tuple[SoquetT, SoquetT]] = [
-            (_reg_to_soq(LeftDangle, reg), in_soqs[reg.name]) for reg in cbloq.signature.lefts()
+            (_reg_to_soq(LeftDangle, reg), cast(SoquetT, in_soqs[reg.name]))
+            for reg in cbloq.signature.lefts()
         ]
 
         for binst, in_soqs, old_out_soqs in cbloq.iter_bloqsoqs():

qualtran/_infra/controlled.py

@@ -49,10 +49,10 @@ def _cvs_convert(
     cvs: Union[
         int,
         np.integer,
-        NDArray[np.number],
+        NDArray[np.integer],
         Sequence[Union[int, np.integer]],
         Sequence[Sequence[Union[int, np.integer]]],
-        Sequence[NDArray[np.number]],
+        Sequence[NDArray[np.integer]],
     ]
 ) -> Tuple[NDArray[np.integer], ...]:
     if isinstance(cvs, (int, np.integer)):
@@ -235,7 +235,7 @@ def from_cirq_cv(
         for qdtype, shape in zip(qdtypes, shapes):
             full_shape = shape + (qdtype.num_qubits,)
             curr_cvs_bits = np.array(cv[idx : idx + int(np.prod(full_shape))]).reshape(full_shape)
-            curr_cvs = np.apply_along_axis(qdtype.from_bits, -1, curr_cvs_bits)
+            curr_cvs = np.apply_along_axis(qdtype.from_bits, -1, curr_cvs_bits)  # type: ignore[arg-type]
             bloq_cvs.append(curr_cvs)
         return CtrlSpec(tuple(qdtypes), tuple(bloq_cvs))
 
@@ -274,7 +274,7 @@ def _get_nice_ctrl_reg_names(reg_names: List[str], n: int) -> Tuple[str, ...]:
         i = 1
     else:
         i = 0
-    names = []
+    names: List[str] = []
     while len(names) < n:
         while True:
             i += 1
@@ -358,13 +358,15 @@ def decompose_bloq(self) -> 'CompositeBloq':
             cbloq = self.subbloq.decompose_bloq()
 
         bb, initial_soqs = BloqBuilder.from_signature(self.signature)
-        ctrl_soqs = [initial_soqs[creg_name] for creg_name in self.ctrl_reg_names]
+        ctrl_soqs: List['SoquetT'] = [initial_soqs[creg_name] for creg_name in self.ctrl_reg_names]
 
         soq_map: List[Tuple[SoquetT, SoquetT]] = []
         for binst, in_soqs, old_out_soqs in cbloq.iter_bloqsoqs():
             in_soqs = bb.map_soqs(in_soqs, soq_map)
             new_bloq, adder = binst.bloq.get_ctrl_system(self.ctrl_spec)
-            ctrl_soqs, new_out_soqs = adder(bb, ctrl_soqs=ctrl_soqs, in_soqs=in_soqs)
+            adder_output = adder(bb, ctrl_soqs=ctrl_soqs, in_soqs=in_soqs)
+            ctrl_soqs = list(adder_output[0])
+            new_out_soqs = adder_output[1]
             soq_map.extend(zip(old_out_soqs, new_out_soqs))
 
         fsoqs = bb.map_soqs(cbloq.final_soqs(), soq_map)
@@ -460,6 +462,7 @@ def as_cirq_op(
         ctrl_regs = {reg_name: cirq_quregs.pop(reg_name) for reg_name in self.ctrl_reg_names}
         ctrl_qubits = [q for reg in ctrl_regs.values() for q in reg.reshape(-1)]
         sub_op, cirq_quregs = self.subbloq.as_cirq_op(qubit_manager, **cirq_quregs)
+        assert sub_op is not None
         return (
             sub_op.controlled_by(*ctrl_qubits, control_values=self.ctrl_spec.to_cirq_cv()),
             cirq_quregs | ctrl_regs,

qualtran/_infra/controlled_test.py

@@ -58,7 +58,7 @@ def test_ctrl_spec():
     cspec1 = CtrlSpec()
     assert cspec1 == CtrlSpec(QBit(), cvs=1)
 
-    cspec2 = CtrlSpec(cvs=np.ones(27).reshape((3, 3, 3)))
+    cspec2 = CtrlSpec(cvs=np.ones(27, dtype=np.intc).reshape((3, 3, 3)))
     assert cspec2.shapes == ((3, 3, 3),)
     assert cspec2 != cspec1
 

qualtran/_infra/gate_with_registers.py

@@ -451,7 +451,7 @@ def controlled(
 
     # pylint: disable=signature-differs
     @overload
-    def controlled(self, ctrl_spec: Optional['CtrlSpec'] = None) -> 'GateWithRegisters':
+    def controlled(self, *, ctrl_spec: Optional['CtrlSpec'] = None) -> 'GateWithRegisters':
         """Bloq-style API to construct a controlled Bloq. See `Bloq.controlled()`."""
 
     def controlled(

qualtran/_infra/gate_with_registers_test.py

@@ -79,7 +79,6 @@ def test_gate_with_registers():
     assert (
         tg.controlled(num_controls=1, control_values=[0])
         == tg.controlled(control_values=[0], control_qid_shape=(2,))
-        == tg.controlled(CtrlSpec(cvs=0))
         == tg.controlled(ctrl_spec=CtrlSpec(cvs=0))
     )
 

qualtran/bloqs/arithmetic/addition.py

@@ -52,6 +52,7 @@
 
     from qualtran.drawing import WireSymbol
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting.symbolic_counting_utils import SymbolicInt
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -272,7 +273,7 @@ class OutOfPlaceAdder(GateWithRegisters, cirq.ArithmeticGate):  # type: ignore[m
         [Halving the cost of quantum addition](https://arxiv.org/abs/1709.06648)
     """
 
-    bitsize: int
+    bitsize: 'SymbolicInt'
     is_adjoint: bool = False
 
     @property
@@ -287,6 +288,8 @@ def signature(self):
         )
 
     def registers(self) -> Sequence[Union[int, Sequence[int]]]:
+        if not isinstance(self.bitsize, int):
+            raise ValueError(f'Symbolic bitsize {self.bitsize} not supported')
         return [2] * self.bitsize, [2] * self.bitsize, [2] * (self.bitsize + 1)
 
     def apply(self, a: int, b: int, c: int) -> Tuple[int, int, int]:
@@ -309,6 +312,8 @@ def short_name(self) -> str:
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
     ) -> cirq.OP_TREE:
+        if not isinstance(self.bitsize, int):
+            raise ValueError(f'Symbolic bitsize {self.bitsize} not supported')
         a, b, c = quregs['a'][::-1], quregs['b'][::-1], quregs['c'][::-1]
         optree: List[List[cirq.Operation]] = [
             [

qualtran/bloqs/arithmetic/comparison.py

@@ -47,6 +47,7 @@
 if TYPE_CHECKING:
     from qualtran import BloqBuilder
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
+    from qualtran.resource_counting.symbolic_counting_utils import SymbolicInt
     from qualtran.simulation.classical_sim import ClassicalValT
 
 
@@ -404,8 +405,8 @@ class LessThanEqual(GateWithRegisters, cirq.ArithmeticGate):  # type: ignore[mis
         https://static-content.springer.com/esm/art%3A10.1038%2Fs41534-018-0071-5/MediaObjects/41534_2018_71_MOESM1_ESM.pdf
     """
 
-    x_bitsize: int
-    y_bitsize: int
+    x_bitsize: 'SymbolicInt'
+    y_bitsize: 'SymbolicInt'
 
     @cached_property
     def signature(self) -> 'Signature':
@@ -414,6 +415,10 @@ def signature(self) -> 'Signature':
         )
 
     def registers(self) -> Sequence[Union[int, Sequence[int]]]:
+        if isinstance(self.x_bitsize, sympy.Expr):
+            raise ValueError(f'Symbolic x bitsize {self.x_bitsize} not allowed')
+        if isinstance(self.y_bitsize, sympy.Expr):
+            raise ValueError(f'Symbolic y bitsize {self.y_bitsize} not allowed')
         return [2] * self.x_bitsize, [2] * self.y_bitsize, [2]
 
     def with_registers(self, *new_registers) -> "LessThanEqual":
@@ -430,6 +435,10 @@ def on_classical_vals(self, *, x: int, y: int, target: int) -> Dict[str, 'Classi
         return {'x': x, 'y': y, 'target': target ^ (x <= y)}
 
     def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo:
+        if isinstance(self.x_bitsize, sympy.Expr):
+            raise ValueError(f'Symbolic x bitsize {self.x_bitsize} not allowed')
+        if isinstance(self.y_bitsize, sympy.Expr):
+            raise ValueError(f'Symbolic y bitsize {self.y_bitsize} not allowed')
         wire_symbols = ["In(x)"] * self.x_bitsize
         wire_symbols += ["In(y)"] * self.y_bitsize
         wire_symbols += ['⨁(x <= y)']
@@ -581,8 +590,8 @@ class GreaterThan(Bloq):
         target: A single bit output register to store the result of A > B.
     """
 
-    a_bitsize: int
-    b_bitsize: int
+    a_bitsize: 'SymbolicInt'
+    b_bitsize: 'SymbolicInt'
 
     @property
     def signature(self):
@@ -690,7 +699,7 @@ def build_composite_bloq(
 
         # Allocate lists to store ancillas generated by the logical-and and control pairs input
         # into logical-ands.
-        ancillas = []
+        ancillas: List[SoquetT] = []
         and_ctrls = []
 
         # If the input registers are unsigned we need to append a sign bit to them in order to use