refactor: Add comprehensive type hints to backend modules

.github/workflows/python-testing.yml

@@ -50,5 +50,8 @@ jobs:
           pip install poetry
           poetry install --extras dev
 
+      - name: Run type checking
+        run: poetry run mypy qumat/
+
       - name: Run tests
         run: poetry run pytest testing/ -v

pyproject.toml

@@ -12,17 +12,41 @@ dependencies = [
     "cirq>=1.5.0,<1.6.0",
     "amazon-braket-sdk>=1.102.6,<2.0",
     "sympy>=1.14.0,<2.0",
+    "flask (>=3.1.2,<4.0.0)",
 ]
 
 [project.optional-dependencies]
-dev = ["pytest>=8.1.1", "ruff>=0.13.1", "pre-commit>=3.0.0"]
+dev = ["pytest>=8.1.1", "ruff>=0.13.1", "pre-commit>=3.0.0", "mypy>=1.0.0"]
 
 [tool.pytest.ini_options]
 testpaths = ["testing"]
 python_files = "test_*.py"
 python_functions = "test_*"
 addopts = ["-v", "--tb=short"]
 
+[tool.mypy]
+python_version = "3.10"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+check_untyped_defs = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+strict_equality = true
+
+# Ignore missing type stubs for third-party quantum libraries
+[[tool.mypy.overrides]]
+module = [
+    "qiskit.*",
+    "qiskit_aer.*",
+    "cirq.*",
+    "braket.*",
+    "sympy.*"
+]
+ignore_missing_imports = true
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

qumat/amazon_braket_backend.py

@@ -14,12 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from typing import Any, cast
+
+import numpy as np
 from braket.aws import AwsDevice
-from braket.devices import LocalSimulator
 from braket.circuits import Circuit, FreeParameter
+from braket.devices import LocalSimulator
 
 
-def initialize_backend(backend_config):
+def initialize_backend(backend_config: dict[str, Any]) -> LocalSimulator | AwsDevice:
     backend_options = backend_config["backend_options"]
     simulator_type = backend_options.get("simulator_type", "default")
     if simulator_type == "local":
@@ -33,55 +36,63 @@ def initialize_backend(backend_config):
         return AwsDevice("arn:aws:braket:::device/quantum-simulator/amazon/sv1")
 
 
-def create_empty_circuit(num_qubits: int | None = None):
+def create_empty_circuit(num_qubits: int | None = None) -> Circuit:
     circuit = Circuit()
     if num_qubits is not None:
         for i in range(num_qubits):
             circuit.i(i)
     return circuit
 
 
-def apply_not_gate(circuit, qubit_index):
+def apply_not_gate(circuit: Circuit, qubit_index: int) -> None:
     circuit.x(qubit_index)
 
 
-def apply_hadamard_gate(circuit, qubit_index):
+def apply_hadamard_gate(circuit: Circuit, qubit_index: int) -> None:
     circuit.h(qubit_index)
 
 
-def apply_cnot_gate(circuit, control_qubit_index, target_qubit_index):
+def apply_cnot_gate(circuit: Circuit, control_qubit_index: int, target_qubit_index: int) -> None:
     circuit.cnot(control_qubit_index, target_qubit_index)
 
 
 def apply_toffoli_gate(
-    circuit, control_qubit_index1, control_qubit_index2, target_qubit_index
-):
+    circuit: Circuit,
+    control_qubit_index1: int,
+    control_qubit_index2: int,
+    target_qubit_index: int,
+) -> None:
     circuit.ccnot(control_qubit_index1, control_qubit_index2, target_qubit_index)
 
 
-def apply_swap_gate(circuit, qubit_index1, qubit_index2):
+def apply_swap_gate(circuit: Circuit, qubit_index1: int, qubit_index2: int) -> None:
     circuit.swap(qubit_index1, qubit_index2)
 
 
 def apply_cswap_gate(
-    circuit, control_qubit_index, target_qubit_index1, target_qubit_index2
-):
+    circuit: Circuit,
+    control_qubit_index: int,
+    target_qubit_index1: int,
+    target_qubit_index2: int,
+) -> None:
     circuit.cswap(control_qubit_index, target_qubit_index1, target_qubit_index2)
 
 
-def apply_pauli_x_gate(circuit, qubit_index):
+def apply_pauli_x_gate(circuit: Circuit, qubit_index: int) -> None:
     circuit.x(qubit_index)
 
 
-def apply_pauli_y_gate(circuit, qubit_index):
+def apply_pauli_y_gate(circuit: Circuit, qubit_index: int) -> None:
     circuit.y(qubit_index)
 
 
-def apply_pauli_z_gate(circuit, qubit_index):
+def apply_pauli_z_gate(circuit: Circuit, qubit_index: int) -> None:
     circuit.z(qubit_index)
 
 
-def execute_circuit(circuit, backend, backend_config):
+def execute_circuit(
+    circuit: Circuit, backend: LocalSimulator | AwsDevice, backend_config: dict[str, Any]
+) -> dict[str, int]:
     shots = backend_config["backend_options"].get("shots", 1)
     parameter_values = backend_config.get("parameter_values", {})
     if parameter_values and circuit.parameters:
@@ -95,57 +106,63 @@ def execute_circuit(circuit, backend, backend_config):
     else:
         task = backend.run(circuit, shots=shots)
     result = task.result()
-    return result.measurement_counts
+    return cast(dict[str, int], result.measurement_counts)
 
 
 # placeholder method for use in the testing suite
-def get_final_state_vector(circuit, backend, backend_config):
+def get_final_state_vector(
+    circuit: Circuit, backend: LocalSimulator | AwsDevice, backend_config: dict[str, Any]
+) -> np.ndarray:
     circuit.state_vector()
     result = backend.run(circuit, shots=0).result()
     state_vector = result.values[0]
 
-    return state_vector
+    return cast(np.ndarray, state_vector)
 
 
-def draw_circuit(circuit):
+def draw_circuit(circuit: Circuit) -> None:
     # Unfortunately, Amazon Braket does not have direct support for drawing circuits in the same way
     # as Qiskit and Cirq. You would typically visualize Amazon Braket circuits using external tools.
     # For simplicity, we'll print the circuit object which gives some textual representation.
     print(circuit)
 
 
-def apply_rx_gate(circuit, qubit_index, angle):
+def apply_rx_gate(circuit: Circuit, qubit_index: int, angle: float | str) -> None:
     if isinstance(angle, (int, float)):
         circuit.rx(qubit_index, angle)
     else:
         param = FreeParameter(angle)
         circuit.rx(qubit_index, param)
 
 
-def apply_ry_gate(circuit, qubit_index, angle):
+def apply_ry_gate(circuit: Circuit, qubit_index: int, angle: float | str) -> None:
     if isinstance(angle, (int, float)):
         circuit.ry(qubit_index, angle)
     else:
         param = FreeParameter(angle)
         circuit.ry(qubit_index, param)
 
 
-def apply_rz_gate(circuit, qubit_index, angle):
+def apply_rz_gate(circuit: Circuit, qubit_index: int, angle: float | str) -> None:
     if isinstance(angle, (int, float)):
         circuit.rz(qubit_index, angle)
     else:
         param = FreeParameter(angle)
         circuit.rz(qubit_index, param)
 
 
-def apply_u_gate(circuit, qubit_index, theta, phi, lambd):
+def apply_u_gate(
+    circuit: Circuit, qubit_index: int, theta: float, phi: float, lambd: float
+) -> None:
     # U(θ, φ, λ) = Rz(φ) · Ry(θ) · Rz(λ)
     circuit.rz(qubit_index, lambd)
     circuit.ry(qubit_index, theta)
     circuit.rz(qubit_index, phi)
 
 
-def calculate_prob_zero(results, ancilla_qubit, num_qubits):
+def calculate_prob_zero(
+    results: dict[str, int] | list[dict[str, int]], ancilla_qubit: int, num_qubits: int
+) -> float:
     """
     Calculate the probability of measuring the ancilla qubit in |0> state.
 

qumat/cirq_backend.py

@@ -14,11 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from typing import Any
+
 import cirq
+import numpy as np
 import sympy
 
 
-def initialize_backend(backend_config):
+def initialize_backend(backend_config: dict[str, Any]) -> cirq.Simulator:
     # Assuming 'simulator_type' specifies the type of simulator in Cirq
     simulator_type = backend_config.get("backend_options", {}).get(
         "simulator_type", "default"
@@ -31,7 +34,7 @@ def initialize_backend(backend_config):
     return cirq.Simulator()
 
 
-def create_empty_circuit(num_qubits: int | None = None):
+def create_empty_circuit(num_qubits: int | None = None) -> cirq.Circuit:
     circuit = cirq.Circuit()
     if num_qubits is not None:
         qubits = [cirq.LineQubit(i) for i in range(num_qubits)]
@@ -40,62 +43,72 @@ def create_empty_circuit(num_qubits: int | None = None):
     return circuit
 
 
-def apply_not_gate(circuit, qubit_index):
+def apply_not_gate(circuit: cirq.Circuit, qubit_index: int) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.X(qubit))
 
 
-def apply_hadamard_gate(circuit, qubit_index):
+def apply_hadamard_gate(circuit: cirq.Circuit, qubit_index: int) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.H(qubit))
 
 
-def apply_cnot_gate(circuit, control_qubit_index, target_qubit_index):
+def apply_cnot_gate(
+    circuit: cirq.Circuit, control_qubit_index: int, target_qubit_index: int
+) -> None:
     control_qubit = cirq.LineQubit(control_qubit_index)
     target_qubit = cirq.LineQubit(target_qubit_index)
     circuit.append(cirq.CNOT(control_qubit, target_qubit))
 
 
 def apply_toffoli_gate(
-    circuit, control_qubit_index1, control_qubit_index2, target_qubit_index
-):
+    circuit: cirq.Circuit,
+    control_qubit_index1: int,
+    control_qubit_index2: int,
+    target_qubit_index: int,
+) -> None:
     control_qubit1 = cirq.LineQubit(control_qubit_index1)
     control_qubit2 = cirq.LineQubit(control_qubit_index2)
     target_qubit = cirq.LineQubit(target_qubit_index)
     circuit.append(cirq.CCX(control_qubit1, control_qubit2, target_qubit))
 
 
-def apply_swap_gate(circuit, qubit_index1, qubit_index2):
+def apply_swap_gate(circuit: cirq.Circuit, qubit_index1: int, qubit_index2: int) -> None:
     qubit1 = cirq.LineQubit(qubit_index1)
     qubit2 = cirq.LineQubit(qubit_index2)
     circuit.append(cirq.SWAP(qubit1, qubit2))
 
 
 def apply_cswap_gate(
-    circuit, control_qubit_index, target_qubit_index1, target_qubit_index2
-):
+    circuit: cirq.Circuit,
+    control_qubit_index: int,
+    target_qubit_index1: int,
+    target_qubit_index2: int,
+) -> None:
     control_qubit = cirq.LineQubit(control_qubit_index)
     target_qubit1 = cirq.LineQubit(target_qubit_index1)
     target_qubit2 = cirq.LineQubit(target_qubit_index2)
     circuit.append(cirq.CSWAP(control_qubit, target_qubit1, target_qubit2))
 
 
-def apply_pauli_x_gate(circuit, qubit_index):
+def apply_pauli_x_gate(circuit: cirq.Circuit, qubit_index: int) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.X(qubit))
 
 
-def apply_pauli_y_gate(circuit, qubit_index):
+def apply_pauli_y_gate(circuit: cirq.Circuit, qubit_index: int) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.Y(qubit))
 
 
-def apply_pauli_z_gate(circuit, qubit_index):
+def apply_pauli_z_gate(circuit: cirq.Circuit, qubit_index: int) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.Z(qubit))
 
 
-def execute_circuit(circuit, backend, backend_config):
+def execute_circuit(
+    circuit: cirq.Circuit, backend: cirq.Simulator, backend_config: dict[str, Any]
+) -> list[dict[int, int]]:
     # handle 0-qubit circuits before adding measurements
     if not circuit.all_qubits():
         shots = backend_config["backend_options"].get("shots", 1)
@@ -122,42 +135,48 @@ def execute_circuit(circuit, backend, backend_config):
         return [result.histogram(key="result")]
 
 
-def draw_circuit(circuit):
+def draw_circuit(circuit: cirq.Circuit) -> None:
     print(circuit)
 
 
-def apply_rx_gate(circuit, qubit_index, angle):
+def apply_rx_gate(circuit: cirq.Circuit, qubit_index: int, angle: float | str) -> None:
     param = sympy.Symbol(angle) if isinstance(angle, str) else angle
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.rx(param).on(qubit))
 
 
-def apply_ry_gate(circuit, qubit_index, angle):
+def apply_ry_gate(circuit: cirq.Circuit, qubit_index: int, angle: float | str) -> None:
     param = sympy.Symbol(angle) if isinstance(angle, str) else angle
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.ry(param).on(qubit))
 
 
-def apply_rz_gate(circuit, qubit_index, angle):
+def apply_rz_gate(circuit: cirq.Circuit, qubit_index: int, angle: float | str) -> None:
     param = sympy.Symbol(angle) if isinstance(angle, str) else angle
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.rz(param).on(qubit))
 
 
-def apply_u_gate(circuit, qubit_index, theta, phi, lambd):
+def apply_u_gate(
+    circuit: cirq.Circuit, qubit_index: int, theta: float, phi: float, lambd: float
+) -> None:
     qubit = cirq.LineQubit(qubit_index)
     circuit.append(cirq.rz(lambd).on(qubit))
     circuit.append(cirq.ry(theta).on(qubit))
     circuit.append(cirq.rz(phi).on(qubit))
 
 
-def get_final_state_vector(circuit, backend, backend_config):
+def get_final_state_vector(
+    circuit: cirq.Circuit, backend: cirq.Simulator, backend_config: dict[str, Any]
+) -> np.ndarray:
     simulator = cirq.Simulator()
     result = simulator.simulate(circuit)
     return result.final_state_vector
 
 
-def calculate_prob_zero(results, ancilla_qubit, num_qubits):
+def calculate_prob_zero(
+    results: list[dict[int, int]] | dict[int, int], ancilla_qubit: int, num_qubits: int
+) -> float:
     """
     Calculate the probability of measuring the ancilla qubit in |0> state.