Mypy issues part 4 (#906)

* Mypy issues part 4

This PR features many fixes to get sympy.Expr to type correctly in
a bunch of places.

Another change was adding iteration_length_or_zero() so we can safely
get iteration_length from dtypes even if they don't have one defined.

Author: dstrain115

dev_tools/conf/mypy.ini

@@ -20,7 +20,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # Non-Google
-[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*]
+[mypy-sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,plotly.*,dash.*,tensorflow_docs.*,fxpmath.*,ipywidgets.*,cachetools.*,pydot.*,nbformat.*,nbconvert.*,openfermion.*]
 follow_imports = silent
 ignore_missing_imports = true
 

qualtran/_infra/adjoint.py

@@ -144,7 +144,7 @@ def decompose_bloq(self) -> 'CompositeBloq':
         return self.subbloq.decompose_bloq().adjoint()
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
     ) -> cirq.OP_TREE:
         if isinstance(self.subbloq, GateWithRegisters):
             return cirq.inverse(self.subbloq.decompose_from_registers(context=context, **quregs))

qualtran/_infra/bloq.py

@@ -212,7 +212,9 @@ def on_classical_vals(self, **vals: 'ClassicalValT') -> Dict[str, 'ClassicalValT
         except NotImplementedError as e:
             raise NotImplementedError(f"{self} does not support classical simulation: {e}") from e
 
-    def call_classically(self, **vals: 'ClassicalValT') -> Tuple['ClassicalValT', ...]:
+    def call_classically(
+        self, **vals: Union['sympy.Symbol', 'ClassicalValT']
+    ) -> Tuple['ClassicalValT', ...]:
         """Call this bloq on classical data.
 
         Bloq users can call this function to apply bloqs to classical data. If you're
@@ -297,7 +299,7 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> Set['BloqCountT']:
     def call_graph(
         self,
         generalizer: Optional[Union['GeneralizerT', Sequence['GeneralizerT']]] = None,
-        keep: Callable[['Bloq'], bool] = None,
+        keep: Optional[Callable[['Bloq'], bool]] = None,
         max_depth: Optional[int] = None,
     ) -> Tuple['nx.DiGraph', Dict['Bloq', Union[int, 'sympy.Expr']]]:
         """Get the bloq call graph and call totals.

qualtran/_infra/composite_bloq.py

@@ -35,6 +35,7 @@
 import attrs
 import networkx as nx
 import numpy as np
+import sympy
 from numpy.typing import NDArray
 
 from .bloq import Bloq, DecomposeTypeError
@@ -45,7 +46,7 @@
 if TYPE_CHECKING:
     import cirq
 
-    from qualtran.cirq_interop import CirqQuregInT, CirqQuregT
+    from qualtran.cirq_interop._cirq_to_bloq import CirqQuregInT, CirqQuregT
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
@@ -447,7 +448,7 @@ def _create_binst_graph(
 
 
 def _binst_to_cxns(
-    binst: BloqInstance, binst_graph: nx.DiGraph
+    binst: Union[BloqInstance, DanglingT], binst_graph: nx.DiGraph
 ) -> Tuple[List[Connection], List[Connection]]:
     """Helper method to extract all predecessor and successor Connections for a binst."""
     pred_cxns: List[Connection] = []
@@ -494,7 +495,7 @@ def _cxn_to_soq_dict(
         assign = get_assign(cxn)
 
         if me.reg.shape:
-            soqdict[me.reg.name][me.idx] = assign
+            soqdict[me.reg.name][me.idx] = assign  # type: ignore[index]
         else:
             soqdict[me.reg.name] = assign
 
@@ -598,9 +599,9 @@ def _reg_to_soq(
     # Annoyingly, this must be a special case.
     # Otherwise, x[i] = thing will nest *array* objects because our ndarray's type is
     # 'object'. This wouldn't happen--for example--with an integer array.
-    soqs = Soquet(binst, reg)
-    available.add(soqs)
-    return soqs
+    soq = Soquet(binst, reg)
+    available.add(soq)
+    return soq
 
 
 def _process_soquets(
@@ -630,7 +631,6 @@ def _process_soquets(
             the incoming, indexed soquet as well as the register and (left-)index it
             has been mapped to.
     """
-
     for reg in registers:
         try:
             # if we want fancy indexing (which we do), we need numpy
@@ -831,7 +831,9 @@ def from_signature(
         initial_soqs: Dict[str, SoquetT] = {}
         for reg in signature:
             if reg.side & Side.LEFT:
-                initial_soqs[reg.name] = bb.add_register_from_dtype(reg)
+                register = bb.add_register_from_dtype(reg)
+                assert register is not None
+                initial_soqs[reg.name] = register
             else:
                 bb.add_register_from_dtype(reg)
 
@@ -866,7 +868,11 @@ def _new_binst_i(self) -> int:
         return i
 
     def _add_cxn(
-        self, binst: BloqInstance, idxed_soq: Soquet, reg: Register, idx: Tuple[int, ...]
+        self,
+        binst: Union[BloqInstance, DanglingT],
+        idxed_soq: Soquet,
+        reg: Register,
+        idx: Tuple[int, ...],
     ) -> None:
         """Helper function to be used as the base for the `func` argument of `_process_soquets`.
 
@@ -1079,7 +1085,7 @@ def _fin(idxed_soq: Soquet, reg: Register, idx: Tuple[int, ...]):
             connections=self._cxns, signature=signature, bloq_instances=self._binsts
         )
 
-    def allocate(self, n: int = 1, dtype: Optional[QDType] = None) -> Soquet:
+    def allocate(self, n: Union[int, sympy.Expr] = 1, dtype: Optional[QDType] = None) -> Soquet:
         from qualtran.bloqs.util_bloqs import Allocate
 
         if dtype is not None:

qualtran/_infra/data_types.py

@@ -89,6 +89,13 @@ def assert_valid_classical_val(self, val: Any, debug_str: str = 'val'):
             debug_str: Optional debugging information to use in exception messages.
         """
 
+    def iteration_length_or_zero(self) -> Union[int, sympy.Expr]:
+        """Safe version of iteration length.
+
+        Returns the iteration_length if the type has it or else zero.
+        """
+        return getattr(self, 'iteration_length', 0)
+
     def assert_valid_classical_val_array(self, val_array: NDArray[Any], debug_str: str = 'val'):
         """Raises an exception if `val_array` is not a valid array of classical values
         for this type.

qualtran/_infra/data_types_test.py

@@ -135,7 +135,7 @@ def test_validation_errs():
         QBit().assert_valid_classical_val(-1)
 
     with pytest.raises(ValueError):
-        QBit().assert_valid_classical_val('|0>')
+        QBit().assert_valid_classical_val('|0>')  # type: ignore[arg-type]
 
     with pytest.raises(ValueError):
         QUInt(3).assert_valid_classical_val(8)

qualtran/_infra/gate_with_registers.py

@@ -14,6 +14,7 @@
 
 import abc
 from typing import (
+    Any,
     Collection,
     Dict,
     Iterable,
@@ -36,7 +37,9 @@
 from qualtran._infra.registers import Register, Side
 
 if TYPE_CHECKING:
-    from qualtran import CtrlSpec
+    import quimb.tensor as qtn
+
+    from qualtran import CtrlSpec, SoquetT
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
 
@@ -355,7 +358,7 @@ def on_registers(
     ) -> cirq.Operation:
         return self.on(*merge_qubits(self.signature, **qubit_regs))
 
-    def __pow__(self, power: int) -> 'GateWithRegisters':
+    def __pow__(self, power: int) -> 'Bloq':
         bloq = self if power > 0 else self.adjoint()
         if abs(power) == 1:
             return bloq

qualtran/_infra/quantum_graph.py

@@ -129,7 +129,7 @@ def _singleton_error(self, x):
     raise ValueError("Do not instantiate a new DanglingT. Use `LeftDangle` or `RightDangle`.")
 
 
-DanglingT.__init__ = _singleton_error
+DanglingT.__init__ = _singleton_error  # type: ignore[method-assign]
 
 
 @frozen

qualtran/bloqs/arithmetic/addition.py

@@ -14,7 +14,7 @@
 import itertools
 import math
 from functools import cached_property
-from typing import Any, Dict, Iterable, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import cirq
 import numpy as np
@@ -253,7 +253,7 @@ def _add_diff_size_regs() -> Add:
 
 
 @frozen
-class OutOfPlaceAdder(GateWithRegisters, cirq.ArithmeticGate):
+class OutOfPlaceAdder(GateWithRegisters, cirq.ArithmeticGate):  # type: ignore[misc]
     r"""An n-bit addition gate.
 
     Implements $U|a\rangle|b\rangle 0\rangle \rightarrow |a\rangle|b\rangle|a+b\rangle$
@@ -310,11 +310,14 @@ def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
     ) -> cirq.OP_TREE:
         a, b, c = quregs['a'][::-1], quregs['b'][::-1], quregs['c'][::-1]
-        optree = [
+        optree: List[List[cirq.Operation]] = [
             [
-                [cirq.CX(a[i], b[i]), cirq.CX(a[i], c[i])],
+                cirq.CX(a[i], b[i]),
+                cirq.CX(a[i], c[i]),
                 And().on(b[i], c[i], c[i + 1]),
-                [cirq.CX(a[i], b[i]), cirq.CX(a[i], c[i + 1]), cirq.CX(b[i], c[i])],
+                cirq.CX(a[i], b[i]),
+                cirq.CX(a[i], c[i + 1]),
+                cirq.CX(b[i], c[i]),
             ]
             for i in range(self.bitsize)
         ]
@@ -418,13 +421,13 @@ def on_classical_vals(
         else:
             return {'x': x + self.k}
 
-        if (self.cvs == ctrls).all():
+        if np.all(self.cvs == ctrls):
             x = x + self.k
 
         return {'ctrls': ctrls, 'x': x}
 
     def build_composite_bloq(
-        self, bb: 'BloqBuilder', x: SoquetT, **regs: SoquetT
+        self, bb: 'BloqBuilder', x: Soquet, **regs: SoquetT
     ) -> Dict[str, 'SoquetT']:
         # Assign registers to variables and allocate ancilla bits for classical integer k.
         if len(self.cvs) > 0:
@@ -444,7 +447,7 @@ def build_composite_bloq(
         # controlled.
         for i in range(self.bitsize):
             if binary_rep[i] == 1:
-                if len(self.cvs) > 0:
+                if len(self.cvs) > 0 and ctrls is not None:
                     ctrls, k_split[i] = bb.add(
                         MultiControlX(cvs=self.cvs), ctrls=ctrls, x=k_split[i]
                     )
@@ -453,14 +456,18 @@ def build_composite_bloq(
 
         # Rejoin the qubits representing k for in-place addition.
         k = bb.join(k_split, dtype=x.reg.dtype)
+        if not isinstance(x.reg.dtype, (QInt, QUInt, QMontgomeryUInt)):
+            raise ValueError(
+                "Only QInt, QUInt and QMontgomerUInt types are supported for composite addition."
+            )
         k, x = bb.add(Add(x.reg.dtype, x.reg.dtype), a=k, b=x)
 
         # Resplit the k qubits in order to undo the original bit flips to go from the binary
         # representation back to the zero state.
         k_split = bb.split(k)
         for i in range(self.bitsize):
             if binary_rep[i] == 1:
-                if len(self.cvs) > 0:
+                if len(self.cvs) > 0 and ctrls is not None:
                     ctrls, k_split[i] = bb.add(
                         MultiControlX(cvs=self.cvs), ctrls=ctrls, x=k_split[i]
                     )
@@ -472,7 +479,7 @@ def build_composite_bloq(
         bb.free(k)
 
         # Return the output registers.
-        if len(self.cvs) > 0:
+        if len(self.cvs) > 0 and ctrls is not None:
             return {'ctrls': ctrls, 'x': x}
         else:
             return {'x': x}
@@ -499,7 +506,7 @@ def _simple_add_k_large() -> SimpleAddConstant:
 
 
 @frozen(auto_attribs=True)
-class AddConstantMod(GateWithRegisters, cirq.ArithmeticGate):
+class AddConstantMod(GateWithRegisters, cirq.ArithmeticGate):  # type: ignore[misc]
     """Applies U(add, M)|x> = |(x + add) % M> if x < M else |x>.
 
     Applies modular addition to input register `|x>` given parameters `mod` and `add_val` s.t.