Add type hints to all remaining top-level files

CHANGELOG.md

@@ -48,6 +48,7 @@ Changelog
 -   Compile to XY gates as well as CZ gates on dummy QVMs (@ecpeterson, gh-1151).
 -   `QAM.write_memory` now accepts either a `Sequence` of values or a single
     value (@tommy-moffat, gh-1114).
+-   Added type hints for all remaining top-level files (@karalekas, gh-1150).
 
 ### Bugfixes
 

Makefile

@@ -16,9 +16,11 @@ check-format:
 # The dream is to one day run mypy on the whole tree. For now, limit checks to known-good files.
 .PHONY: check-types
 check-types:
-	mypy pyquil/gates.py pyquil/gate_matrices.py pyquil/noise.py pyquil/numpy_simulator.py \
+	mypy pyquil/gate_matrices.py pyquil/gates.py pyquil/noise.py pyquil/numpy_simulator.py \
+		pyquil/operator_estimation.py pyquil/parser.py pyquil/paulis.py pyquil/pyqvm.py \
 		pyquil/quil.py pyquil/quilatom.py pyquil/quilbase.py pyquil/reference_simulator.py \
-		pyquil/unitary_tools.py pyquil/latex pyquil/simulation pyquil/experiment pyquil/device
+		pyquil/unitary_tools.py pyquil/version.py pyquil/wavefunction.py \
+		pyquil/device pyquil/experiment pyquil/latex pyquil/simulation
 
 .PHONY: check-style
 check-style:

mypy.ini

@@ -19,14 +19,22 @@ warn_unused_ignores = True
 warn_return_any = True
 no_implicit_reexport = True
 
-# Ignore errors in generated parser files
-[mypy-pyquil._parser.gen3.*]
+# Ignore errors in all parser-related files
+[mypy-pyquil._parser.*]
 ignore_errors = True
 
 # Ignore errors in vendored third-party libraries
 [mypy-pyquil.external.*]
 ignore_errors = True
 
-# Ignore errors in test files
+# Ignore errors in all test files
 [mypy-*/tests/*]
 ignore_errors = True
+
+# Ignore errors in the conftest.py file
+[mypy-conftest]
+ignore_errors = True
+
+# Ignore errors in the pyquil/magic.py file
+[mypy-pyquil.magic]
+ignore_errors = True

pyquil/operator_estimation.py

@@ -1,53 +1,53 @@
 import logging
 import warnings
 from math import pi
-from typing import Callable, Dict, List, Union, Tuple, Optional
+from typing import Callable, Generator, List, Mapping, Union, Tuple, Optional, cast
 
 import numpy as np
 
-from pyquil import Program
 from pyquil.api import QuantumComputer
+from pyquil.quil import Program
+from pyquil.quilatom import QubitDesignator
 
 # import the full public API of the pyquil experiment module
-from pyquil.experiment import (
+from pyquil.experiment._group import (
+    _max_weight_state,
+    _max_weight_operator,
+    construct_tpb_graph,
+    group_settings as group_experiments,
+    group_settings_clique_removal as group_experiments_clique_removal,
+    group_settings_greedy as group_experiments_greedy,
+)
+from pyquil.experiment._main import (
+    OperatorEncoder,
+    TomographyExperiment,
+)
+from pyquil.experiment._result import ExperimentResult, ratio_variance
+from pyquil.experiment._setting import (
     _OneQState,
     _pauli_to_product_state,
-    ExperimentResult,
     ExperimentSetting,
-    OperatorEncoder,
     SIC0,
     SIC1,
     SIC2,
     SIC3,
-    SymmetrizationLevel,
-    TomographyExperiment,
     TensorProductState,
     minusX,
     minusY,
     minusZ,
     plusX,
     plusY,
     plusZ,
-    read_json,
-    to_json,
     zeros_state,
 )
-from pyquil.experiment._group import (
-    _max_weight_state,
-    _max_weight_operator,
-    construct_tpb_graph,
-    group_settings as group_experiments,
-    group_settings_clique_removal as group_experiments_clique_removal,
-    group_settings_greedy as group_experiments_greedy,
-)
-from pyquil.experiment._result import ratio_variance
+from pyquil.experiment._symmetrization import SymmetrizationLevel
 from pyquil.gates import RESET, RX, RY, RZ, X
 from pyquil.paulis import is_identity
 
 log = logging.getLogger(__name__)
 
 
-def _one_q_sic_prep(index, qubit):
+def _one_q_sic_prep(index: int, qubit: QubitDesignator) -> Program:
     """Prepare the index-th SIC basis state."""
     if index == 0:
         return Program()
@@ -67,7 +67,7 @@ def _one_q_sic_prep(index, qubit):
     raise ValueError(f"Bad SIC index: {index}")
 
 
-def _one_q_pauli_prep(label, index, qubit):
+def _one_q_pauli_prep(label: str, index: int, qubit: QubitDesignator) -> Program:
     """Prepare the index-th eigenstate of the pauli operator given by label."""
     if index not in [0, 1]:
         raise ValueError(f"Bad Pauli index: {index}")
@@ -93,7 +93,7 @@ def _one_q_pauli_prep(label, index, qubit):
     raise ValueError(f"Bad Pauli label: {label}")
 
 
-def _one_q_state_prep(oneq_state: _OneQState):
+def _one_q_state_prep(oneq_state: _OneQState) -> Program:
     """Prepare a one qubit state.
 
     Either SIC[0-3], X[0-1], Y[0-1], or Z[0-1].
@@ -107,7 +107,7 @@ def _one_q_state_prep(oneq_state: _OneQState):
         raise ValueError(f"Bad state label: {label}")
 
 
-def _local_pauli_eig_meas(op, idx):
+def _local_pauli_eig_meas(op: str, idx: QubitDesignator) -> Program:
     """
     Generate gate sequence to measure in the eigenbasis of a Pauli operator, assuming
     we are only able to measure in the Z eigenbasis. (Note: The unitary operations of this
@@ -175,11 +175,12 @@ def _generate_experiment_programs(
                 "so that groups of parallel settings have compatible observables."
             )
         for qubit, op_str in max_weight_out_op:
+            assert isinstance(qubit, int)
             total_prog += _local_pauli_eig_meas(op_str, qubit)
 
         programs.append(total_prog)
 
-        meas_qubits.append(max_weight_out_op.get_qubits())
+        meas_qubits.append(cast(List[int], max_weight_out_op.get_qubits()))
     return programs, meas_qubits
 
 
@@ -192,7 +193,7 @@ def measure_observables(
     symmetrize_readout: Optional[Union[int, str]] = "None",
     calibrate_readout: Optional[str] = "plus-eig",
     readout_symmetrize: Optional[str] = None,
-):
+) -> Generator[ExperimentResult, None, None]:
     """
     Measure all the observables in a TomographyExperiment.
 
@@ -305,7 +306,8 @@ def measure_observables(
         # either the first column, second column, or both and multiplying along the row.
         for setting in settings:
             # Get the term's coefficient so we can multiply it in later.
-            coeff = complex(setting.out_operator.coefficient)
+            coeff = setting.out_operator.coefficient
+            assert isinstance(coeff, complex)
             if not np.isclose(coeff.imag, 0):
                 raise ValueError(f"{setting}'s out_operator has a complex coefficient.")
             coeff = coeff.real
@@ -328,7 +330,7 @@ def measure_observables(
                 # Obtain calibration program
                 calibr_prog = _calibration_program(qc, tomo_experiment, setting)
                 calibr_qubs = setting.out_operator.get_qubits()
-                calibr_qub_dict = {q: idx for idx, q in enumerate(calibr_qubs)}
+                calibr_qub_dict = {cast(int, q): idx for idx, q in enumerate(calibr_qubs)}
 
                 # Perform symmetrization on the calibration program
                 calibr_results = qc.run_symmetrized_readout(
@@ -389,11 +391,11 @@ def _ops_bool_to_prog(ops_bool: Tuple[bool], qubits: List[int]) -> Program:
 
 def _stats_from_measurements(
     bs_results: np.ndarray,
-    qubit_index_map: Dict,
+    qubit_index_map: Mapping[int, int],
     setting: ExperimentSetting,
     n_shots: int,
     coeff: float = 1.0,
-) -> Tuple[float, float]:
+) -> Tuple[np.ndarray, np.ndarray]:
     """
     :param bs_results: results from running `qc.run`
     :param qubit_index_map: dict mapping qubit to classical register index
@@ -403,7 +405,7 @@ def _stats_from_measurements(
     :return: tuple specifying (mean, variance)
     """
     # Identify classical register indices to select
-    idxs = [qubit_index_map[q] for q, _ in setting.out_operator]
+    idxs = [qubit_index_map[cast(int, q)] for q, _ in setting.out_operator]
     # Pick columns corresponding to qubits with a non-identity out_operation
     obs_strings = bs_results[:, idxs]
     # Transform bits to eigenvalues; ie (+1, -1)
@@ -444,9 +446,11 @@ def _calibration_program(
     calibr_prog += kraus_instructions
     # Prepare the +1 eigenstate for the out operator
     for q, op in setting.out_operator.operations_as_set():
+        assert isinstance(q, int)
         calibr_prog += _one_q_pauli_prep(label=op, index=0, qubit=q)
     # Measure the out operator in this state
     for q, op in setting.out_operator.operations_as_set():
+        assert isinstance(q, int)
         calibr_prog += _local_pauli_eig_meas(op, q)
 
     return calibr_prog

pyquil/parser.py

@@ -16,12 +16,14 @@
 """
 Module for parsing Quil programs from text into PyQuil objects
 """
-from pyquil.quil import Program
+from typing import List
 
 from pyquil._parser.PyQuilListener import run_parser
+from pyquil.quil import Program
+from pyquil.quilbase import AbstractInstruction
 
 
-def parse_program(quil):
+def parse_program(quil: str) -> Program:
     """
     Parse a raw Quil program and return a PyQuil program.
 
@@ -31,7 +33,7 @@ def parse_program(quil):
     return Program(parse(quil))
 
 
-def parse(quil):
+def parse(quil: str) -> List[AbstractInstruction]:
     """
     Parse a raw Quil program and return a corresponding list of PyQuil objects.