Repository: qBraid/qbraid-qir (language: python)

Issue:
Add static type checking with `mypy`
Make qBraid-QIR repository pass checks from [mypy](https://github.com/python/mypy) static type checker.

- Run `mypy` and fix type errors.
- Add any necessary configurations to a `mypy.ini` file in the project root.
- Add `disable_error_code` patterns to `mypy.ini` as necessary, but sparingly.
- Add `# type: ignore` in-line skips as necessary, but sparingly.
- Add `mypy` type check to `tox.ini` [testenv:format-check](https://github.com/qBraid/qbraid-qir/blob/6d25a7f76241fad82927469f3fa5947bb9e31803/tox.ini#L71).


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml
index a9b7673..36e5308 100644
--- a/.github/workflows/format.yml
+++ b/.github/workflows/format.yml
@@ -19,6 +19,6 @@ jobs:
         run: |
           python3 -m pip install --upgrade pip
           python3 -m pip install tox>=4.2.0
-      - name: Check isort, black, headers
+      - name: Check isort, black, mypy, headers
         run: |
           tox -e format-check
\ No newline at end of file
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 31e7ca7..4a812f5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,8 @@ Types of changes:
 
 ### 🌟  Improvements 
 * Re-factor the `BasicQasmVisitor` and improve modularity ( [#142](https://github.com/qBraid/qbraid-qir/pull/142) )
+* Add static type checking with `mypy` ( [#150](https://github.com/qBraid/qbraid-qir/pull/150) )
+* Improve measurement statement parsing logic and add support for range definition and discrete set ( [#150](https://github.com/qBraid/qbraid-qir/pull/150) )
 
 ### 📜  Documentation 
 * Housekeeping updates for release ( [#135](https://github.com/qBraid/qbraid-qir/pull/135) )
diff --git a/mypy.ini b/mypy.ini
new file mode 100644
index 0000000..b973238
--- /dev/null
+++ b/mypy.ini
@@ -0,0 +1,16 @@
+[mypy]
+
+# Ignore missing imports
+ignore_missing_imports = True
+
+# Enable incremental mode
+incremental = True
+
+# Show error codes in output
+show_error_codes = True
+
+# Follow imports for type checking
+follow_imports = normal
+
+# Enable cache
+cache_dir = .mypy_cache
diff --git a/qbraid_qir/cirq/convert.py b/qbraid_qir/cirq/convert.py
index 684f811..31c29b5 100644
--- a/qbraid_qir/cirq/convert.py
+++ b/qbraid_qir/cirq/convert.py
@@ -64,7 +64,7 @@ def cirq_to_qir(circuit: cirq.Circuit, name: Optional[str] = None, **kwargs) ->
     visitor = BasicCirqVisitor(**kwargs)
     module.accept(visitor)
 
-    err = llvm_module.verify()
-    if err is not None:
-        raise CirqConversionError(err)
+    error = llvm_module.verify()
+    if error is not None:
+        raise CirqConversionError(error)
     return llvm_module
diff --git a/qbraid_qir/cirq/elements.py b/qbraid_qir/cirq/elements.py
index 360d6f2..8ca45b7 100644
--- a/qbraid_qir/cirq/elements.py
+++ b/qbraid_qir/cirq/elements.py
@@ -14,7 +14,7 @@
 
 import hashlib
 from abc import ABCMeta, abstractmethod
-from typing import FrozenSet, List, Optional
+from typing import Optional
 
 import cirq
 from pyqir import Context, Module
@@ -50,7 +50,7 @@ def accept(self, visitor):
 
 
 class _Register(_CircuitElement):
-    def __init__(self, register: FrozenSet[cirq.Qid]):
+    def __init__(self, register: list[cirq.Qid]):
         self._register = register
 
     def accept(self, visitor):
@@ -77,7 +77,7 @@ class CirqModule:
         name (str): Name of the module.
         module (Module): QIR Module instance.
         num_qubits (int): Number of qubits in the circuit.
-        elements (List[_CircuitElement]): List of circuit elements.
+        elements (list[_CircuitElement]): list of circuit elements.
 
     Example:
         >>> circuit = cirq.Circuit()
@@ -90,7 +90,7 @@ def __init__(
         name: str,
         module: Module,
         num_qubits: int,
-        elements: List[_CircuitElement],
+        elements: list[_CircuitElement],
     ):
         self._name = name
         self._module = module
@@ -122,7 +122,7 @@ def num_clbits(self) -> int:
     def from_circuit(cls, circuit: cirq.Circuit, module: Optional[Module] = None) -> "CirqModule":
         """Class method. Constructs a CirqModule from a given cirq.Circuit object
         and an optional QIR Module."""
-        elements: List[_CircuitElement] = []
+        elements: list[_CircuitElement] = []
 
         # Register(s). Tentatively using cirq.Qid as input. Better approaches might exist tbd.
         elements.append(_Register(list(circuit.all_qubits())))
diff --git a/qbraid_qir/cirq/opsets.py b/qbraid_qir/cirq/opsets.py
index 7b2bbfe..316c2fd 100644
--- a/qbraid_qir/cirq/opsets.py
+++ b/qbraid_qir/cirq/opsets.py
@@ -12,7 +12,7 @@
 Module mapping supported Cirq gates/operations to pyqir functions.
 
 """
-from typing import Callable, Tuple
+from typing import Callable
 
 import cirq
 import pyqir._native
@@ -38,7 +38,7 @@ def measure_y(builder, qubit, result):
     pyqir._native.mz(builder, qubit, result)
 
 
-PYQIR_OP_MAP = {
+PYQIR_OP_MAP: dict[str, Callable] = {
     # Identity Gate
     "I": i,
     # Single-Qubit Clifford Gates
@@ -71,7 +71,7 @@ def measure_y(builder, qubit, result):
 
 def map_cirq_op_to_pyqir_callable(
     operation: cirq.Operation,
-) -> Tuple[Callable, str]:
+) -> tuple[Callable, str]:
     """
     Maps a Cirq operation to its corresponding PyQIR callable function.
 
@@ -79,7 +79,7 @@ def map_cirq_op_to_pyqir_callable(
         operation (cirq.Operation): The Cirq operation to map.
 
     Returns:
-        Tuple[Callable, str]: Tuple containing the corresponding PyQIR callable function,
+        tuple[Callable, str]: tuple containing the corresponding PyQIR callable function,
                                and a string representing the gate/operation type.
 
     Raises:
diff --git a/qbraid_qir/cirq/passes.py b/qbraid_qir/cirq/passes.py
index baad1e2..e7107b9 100644
--- a/qbraid_qir/cirq/passes.py
+++ b/qbraid_qir/cirq/passes.py
@@ -13,7 +13,7 @@
 
 """
 import itertools
-from typing import List
+from typing import Iterable
 
 import cirq
 
@@ -21,15 +21,15 @@
 from .opsets import map_cirq_op_to_pyqir_callable
 
 
-def _decompose_gate_op(operation: cirq.GateOperation) -> List[cirq.OP_TREE]:
+def _decompose_gate_op(operation: cirq.Operation) -> Iterable[cirq.OP_TREE]:
     """Decomposes a single Cirq gate operation into a sequence of operations
     that are directly supported by PyQIR.
 
     Args:
-        operation (cirq.GateOperation): The gate operation to decompose.
+        operation (cirq.Operation): The gate operation to decompose.
 
     Returns:
-        List[cirq.OP_TREE]: A list of decomposed gate operations.
+        Iterable[cirq.OP_TREE]: A list of decomposed gate operations.
     """
     try:
         # Try converting to PyQIR. If successful, keep the operation.
@@ -58,7 +58,7 @@ def _decompose_unsupported_gates(circuit: cirq.Circuit) -> cirq.Circuit:
         new_ops = []
         for operation in moment:
             if isinstance(operation, cirq.GateOperation):
-                decomposed_ops = _decompose_gate_op(operation)
+                decomposed_ops = list(_decompose_gate_op(operation))
                 new_ops.extend(decomposed_ops)
             elif isinstance(operation, cirq.ClassicallyControlledOperation):
                 new_ops.append(operation)
diff --git a/qbraid_qir/cirq/visitor.py b/qbraid_qir/cirq/visitor.py
index 20408da..79cf518 100644
--- a/qbraid_qir/cirq/visitor.py
+++ b/qbraid_qir/cirq/visitor.py
@@ -14,7 +14,6 @@
 """
 import logging
 from abc import ABCMeta, abstractmethod
-from typing import List
 
 import cirq
 import pyqir
@@ -49,11 +48,11 @@ class BasicCirqVisitor(CircuitElementVisitor):
     """
 
     def __init__(self, initialize_runtime: bool = True, record_output: bool = True):
-        self._module = None
-        self._builder = None
-        self._entry_point = None
-        self._qubit_labels = {}
-        self._measured_qubits = {}
+        self._module: pyqir.Module
+        self._builder: pyqir.Builder
+        self._entry_point: str
+        self._qubit_labels: dict[cirq.Qid, int] = {}
+        self._measured_qubits: dict = {}
         self._initialize_runtime = initialize_runtime
         self._record_output = record_output
 
@@ -89,12 +88,9 @@ def record_output(self, module: CirqModule) -> None:
             result_ref = pyqir.result(self._module.context, i)
             pyqir.rt.result_record_output(self._builder, result_ref, Constant.null(i8p))
 
-    def visit_register(self, qids: List[cirq.Qid]) -> None:
+    def visit_register(self, qids: list[cirq.Qid]) -> None:
         logger.debug("Visiting qids '%s'", str(qids))
 
-        if not isinstance(qids, list):
-            raise TypeError("Parameter is not a list.")
-
         if not all(isinstance(x, cirq.Qid) for x in qids):
             raise TypeError("All elements in the list must be of type cirq.Qid.")
 
@@ -102,7 +98,8 @@ def visit_register(self, qids: List[cirq.Qid]) -> None:
         logger.debug("Added labels for qubits %s", str(qids))
 
     def visit_operation(self, operation: cirq.Operation) -> None:
-        qlabels = [self._qubit_labels.get(bit) for bit in operation.qubits]
+        qlabels = [self._qubit_labels[bit] for bit in operation.qubits]
+
         qubits = [pyqir.qubit(self._module.context, n) for n in qlabels]
         results = [pyqir.result(self._module.context, n) for n in qlabels]
 
@@ -125,7 +122,9 @@ def handle_measurement(pyqir_func):
             # pylint: disable=unnecessary-lambda-assignment
             if op_str in ["Rx", "Ry", "Rz"]:
                 pyqir_func = lambda: temp_pyqir_func(
-                    self._builder, operation._sub_operation.gate._rads, *qubits
+                    self._builder,
+                    operation._sub_operation.gate._rads,  # type: ignore[union-attr]
+                    *qubits,
                 )
             else:
                 pyqir_func = lambda: temp_pyqir_func(self._builder, *qubits)
@@ -149,12 +148,9 @@ def _branch(conds, pyqir_func):
             if op_str.startswith("measure"):
                 handle_measurement(pyqir_func)
             elif op_str in ["Rx", "Ry", "Rz"]:
-                pyqir_func(self._builder, operation.gate._rads, *qubits)
+                pyqir_func(self._builder, operation.gate._rads, *qubits)  # type: ignore[union-attr]
             else:
                 pyqir_func(self._builder, *qubits)
 
     def ir(self) -> str:
         return str(self._module)
-
-    def bitcode(self) -> bytes:
-        return self._module.bitcode()
diff --git a/qbraid_qir/qasm3/analyzer.py b/qbraid_qir/qasm3/analyzer.py
index 7077c00..6da5781 100644
--- a/qbraid_qir/qasm3/analyzer.py
+++ b/qbraid_qir/qasm3/analyzer.py
@@ -12,10 +12,12 @@
 Module with analysis functions for QASM3 visitor
 
 """
-from typing import Any
+from typing import Any, Optional, Union
 
 from openqasm3.ast import (
     BinaryExpression,
+    DiscreteSet,
+    Expression,
     IndexExpression,
     IntegerLiteral,
     RangeDefinition,
@@ -30,7 +32,7 @@ class Qasm3Analyzer:
     """Class with utility functions for analyzing QASM3 elements"""
 
     @staticmethod
-    def analyze_classical_indices(indices: list[IntegerLiteral], var: Variable) -> None:
+    def analyze_classical_indices(indices: list[Any], var: Variable) -> list:
         """Validate the indices for a classical variable.
 
         Args:
@@ -45,18 +47,21 @@ def analyze_classical_indices(indices: list[IntegerLiteral], var: Variable) -> N
         """
         indices_list = []
         var_name = var.name
-        var_dimensions = var.dims
+        var_dimensions: Optional[list[int]] = var.dims
 
-        if not var_dimensions:
+        if var_dimensions is None or len(var_dimensions) == 0:
             raise_qasm3_error(
                 message=f"Indexing error. Variable {var_name} is not an array",
                 err_type=Qasm3ConversionError,
                 span=indices[0].span,
             )
-        if len(indices) != len(var_dimensions):
+        if isinstance(indices, DiscreteSet):
+            indices = indices.values
+
+        if len(indices) != len(var_dimensions):  # type: ignore[arg-type]
             raise_qasm3_error(
                 message=f"Invalid number of indices for variable {var_name}. "
-                f"Expected {len(var_dimensions)} but got {len(indices)}",
+                f"Expected {len(var_dimensions)} but got {len(indices)}",  # type: ignore[arg-type]
                 err_type=Qasm3ConversionError,
                 span=indices[0].span,
             )
@@ -78,7 +83,7 @@ def analyze_classical_indices(indices: list[IntegerLiteral], var: Variable) -> N
                     span=index.span,
                 )
             index_value = index.value
-            curr_dimension = var_dimensions[i]
+            curr_dimension = var_dimensions[i]  # type: ignore[index]
 
             if index_value < 0 or index_value >= curr_dimension:
                 raise_qasm3_error(
@@ -92,29 +97,35 @@ def analyze_classical_indices(indices: list[IntegerLiteral], var: Variable) -> N
         return indices_list
 
     @staticmethod
-    def analyze_index_expression(index_expr: IndexExpression) -> tuple[str, list[list]]:
+    def analyze_index_expression(
+        index_expr: IndexExpression,
+    ) -> tuple[str, list[Union[Any, Expression, RangeDefinition]]]:
         """analyze an index expression to get the variable name and indices.
 
         Args:
             index_expr (IndexExpression): The index expression to analyze.
 
         Returns:
-            tuple[str, list[list]]: The variable name and indices.
+            tuple[str, list[Any]]: The variable name and indices.
 
         """
-        indices = []
-        var_name = None
+        indices: list[Any] = []
+        var_name = ""
         comma_separated = False
 
         if isinstance(index_expr.collection, IndexExpression):
             while isinstance(index_expr, IndexExpression):
-                indices.append(index_expr.index[0])
-                index_expr = index_expr.collection
+                if isinstance(index_expr.index, list):
+                    indices.append(index_expr.index[0])
+                    index_expr = index_expr.collection
         else:
             comma_separated = True
-            indices = index_expr.index
-
-        var_name = index_expr.collection.name if comma_separated else index_expr.name
+            indices = index_expr.index  # type: ignore[assignment]
+        var_name = (
+            index_expr.collection.name  # type: ignore[attr-defined]
+            if comma_separated
+            else index_expr.name  # type: ignore[attr-defined]
+        )
         if not comma_separated:
             indices = indices[::-1]
 
@@ -169,7 +180,7 @@ def analyse_branch_condition(condition: Any) -> bool:
                     err_type=Qasm3ConversionError,
                     span=condition.span,
                 )
-            return condition.rhs.value != 0
+            return condition.rhs.value != 0  # type: ignore[attr-defined]
         if not isinstance(condition, IndexExpression):
             raise_qasm3_error(
                 message=(
diff --git a/qbraid_qir/qasm3/elements.py b/qbraid_qir/qasm3/elements.py
index 4aee167..775aef9 100644
--- a/qbraid_qir/qasm3/elements.py
+++ b/qbraid_qir/qasm3/elements.py
@@ -57,7 +57,7 @@ class Variable:
         name (str): Name of the variable.
         base_type (Any): Base type of the variable.
         base_size (int): Base size of the variable.
-        dims (List[int]): Dimensions of the variable.
+        dims (list[int]): Dimensions of the variable.
         value (Optional[Union[int, float, list]]): Value of the variable.
         is_constant (bool): Flag indicating if the variable is constant.
 
@@ -122,7 +122,7 @@ class Qasm3Module:
         module (Module): QIR Module instance.
         num_qubits (int): Number of qubits in the circuit.
         num_clbits (int): Number of classical bits in the circuit.
-        elements (List[Statement]): List of openqasm3 Statements.
+        elements (list[Statement]): list of openqasm3 Statements.
     """
 
     # pylint: disable-next=too-many-arguments
@@ -166,20 +166,24 @@ def from_program(cls, program: Program, module: Optional[Module] = None):
         Class method. Construct a Qasm3Module from a given openqasm3.ast.Program object
         and an optional QIR Module.
         """
-        elements = []
+        elements: list[Union[_Register, _Statement]] = []
 
         num_qubits = 0
         num_clbits = 0
         for statement in program.statements:
             if isinstance(statement, QubitDeclaration):
-                size = 1 if statement.size is None else statement.size.value
+                size = 1
+                if statement.size:
+                    size = statement.size.value  # type: ignore[attr-defined]
                 num_qubits += size
                 elements.append(_Register(statement))
 
             elif isinstance(statement, ClassicalDeclaration) and isinstance(
                 statement.type, BitType
             ):
-                size = 1 if statement.type.size is None else statement.type.size.value
+                size = 1
+                if statement.type.size:
+                    size = statement.type.size.value  # type: ignore[attr-defined]
                 num_clbits += size
                 elements.append(_Register(statement))
                 # as bit arrays are just 0 / 1 values, we can treat them as
@@ -191,7 +195,7 @@ def from_program(cls, program: Program, module: Optional[Module] = None):
 
         if module is None:
             # pylint: disable-next=too-many-function-args
-            module = Module(qirContext(), generate_module_id(program))
+            module = Module(qirContext(), generate_module_id())
 
         return cls(
             name="main",
diff --git a/qbraid_qir/qasm3/expressions.py b/qbraid_qir/qasm3/expressions.py
index bbe8382..e73eb5c 100644
--- a/qbraid_qir/qasm3/expressions.py
+++ b/qbraid_qir/qasm3/expressions.py
@@ -37,7 +37,7 @@ class Qasm3ExprEvaluator:
     visitor_obj = None
 
     @classmethod
-    def set_visitor_obj(cls, visitor_obj):
+    def set_visitor_obj(cls, visitor_obj) -> None:
         cls.visitor_obj = visitor_obj
 
     @classmethod
@@ -171,7 +171,7 @@ def evaluate_expression(cls, expression, const_expr: bool = False, reqd_type=Non
                 expression.span,
             )
 
-        def _process_variable(var_name, indices=None):
+        def _process_variable(var_name: str, indices=None):
             cls._check_var_in_scope(var_name, expression)
             cls._check_var_constant(var_name, const_expr, expression)
             cls._check_var_type(var_name, reqd_type, expression)
@@ -200,15 +200,18 @@ def _process_variable(var_name, indices=None):
             target = expression.target
             index = expression.index
 
-            if not isinstance(target, Identifier):
+            if isinstance(target, Identifier):
+                var_name = target.name
+                cls._check_var_in_scope(var_name, expression)
+                dimensions = cls.visitor_obj._get_from_visible_scope(  # type: ignore[union-attr]
+                    var_name
+                ).dims
+            else:
                 raise_qasm3_error(
                     message=f"Unsupported target type {type(target)} for sizeof expression",
                     err_type=Qasm3ConversionError,
                     span=expression.span,
                 )
-            var_name = target.name
-            cls._check_var_in_scope(var_name, expression)
-            dimensions = cls.visitor_obj._get_from_visible_scope(var_name).dims
 
             if dimensions is None or len(dimensions) == 0:
                 raise_qasm3_error(
@@ -222,7 +225,7 @@ def _process_variable(var_name, indices=None):
                 return dimensions[0]
 
             index = cls.evaluate_expression(index, const_expr, reqd_type=Qasm3IntType)
-
+            assert index is not None and isinstance(index, int)
             if index < 0 or index >= len(dimensions):
                 raise_qasm3_error(
                     f"Index {index} out of bounds for array {var_name} with "
@@ -268,7 +271,7 @@ def _process_variable(var_name, indices=None):
             # function will not return a reqd / const type
             # Reference : https://openqasm.com/language/types.html#compile-time-constants
             # para      : 5
-            return cls.visitor_obj._visit_function_call(expression)
+            return cls.visitor_obj._visit_function_call(expression)  # type: ignore[union-attr]
 
         raise_qasm3_error(
             f"Unsupported expression type {type(expression)}", Qasm3ConversionError, expression.span
diff --git a/qbraid_qir/qasm3/maps.py b/qbraid_qir/qasm3/maps.py
index ae9de26..3aabc4b 100644
--- a/qbraid_qir/qasm3/maps.py
+++ b/qbraid_qir/qasm3/maps.py
@@ -14,7 +14,7 @@
 """
 
 
-from typing import Union
+from typing import Callable, Union
 
 import numpy as np
 import pyqir
@@ -36,7 +36,13 @@
 from .exceptions import Qasm3ConversionError
 from .linalg import kak_decomposition_angles
 
-OPERATOR_MAP = {
+# Define the type for the operator functions
+OperatorFunction = Union[
+    Callable[[Union[int, float, bool]], Union[int, float, bool]],
+    Callable[[Union[int, float, bool], Union[int, float, bool]], Union[int, float, bool]],
+]
+
+OPERATOR_MAP: dict[str, OperatorFunction] = {
     "+": lambda x, y: x + y,
     "-": lambda x, y: x - y,
     "*": lambda x, y: x * y,
@@ -61,7 +67,7 @@
 }
 
 
-def qasm3_expression_op_map(op_name: str, *args):
+def qasm3_expression_op_map(op_name: str, *args) -> Union[float, int, bool]:
     """
     Return the result of applying the given operator to the given operands.
 
@@ -75,7 +81,8 @@ def qasm3_expression_op_map(op_name: str, *args):
         (Union[float, int, bool]): The result of applying the operator to the operands.
     """
     try:
-        return OPERATOR_MAP[op_name](*args)
+        operator = OPERATOR_MAP[op_name]
+        return operator(*args)
     except KeyError as exc:
         raise Qasm3ConversionError(f"Unsupported / undeclared QASM operator: {op_name}") from exc
 
diff --git a/qbraid_qir/qasm3/transformer.py b/qbraid_qir/qasm3/transformer.py
index e5ad65b..0d04131 100644
--- a/qbraid_qir/qasm3/transformer.py
+++ b/qbraid_qir/qasm3/transformer.py
@@ -12,12 +12,12 @@
 Module with transformation functions for QASM3 visitor
 
 """
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 from openqasm3.ast import (
     BinaryExpression,
     DiscreteSet,
-    FloatLiteral,
+    Expression,
     Identifier,
     IndexedIdentifier,
     IndexExpression,
@@ -33,6 +33,8 @@
 from .expressions import Qasm3ExprEvaluator
 from .validator import Qasm3Validator
 
+# mypy: disable-error-code="attr-defined, union-attr"
+
 
 class Qasm3Transformer:
     """Class with utility functions for transforming QASM3 elements"""
@@ -40,7 +42,7 @@ class Qasm3Transformer:
     visitor_obj = None
 
     @classmethod
-    def set_visitor_obj(cls, visitor_obj):
+    def set_visitor_obj(cls, visitor_obj) -> None:
         cls.visitor_obj = visitor_obj
 
     @staticmethod
@@ -128,19 +130,18 @@ def transform_gate_qubits(
                     f"Indexing '{qubit.name.name}' not supported in gate definition",
                     span=qubit.span,
                 )
-
-            gate_op.qubits[i] = qubit_map[qubit.name]
+            gate_qubit_name = qubit.name
+            if isinstance(gate_qubit_name, Identifier):
+                gate_qubit_name = gate_qubit_name.name
+            gate_op.qubits[i] = qubit_map[gate_qubit_name]
 
     @staticmethod
-    def transform_gate_params(
-        gate_op: QuantumGate, param_map: dict[str, Union[FloatLiteral, IntegerLiteral]]
-    ) -> None:
+    def transform_gate_params(gate_op: QuantumGate, param_map: dict[str, Expression]) -> None:
         """Transform the parameters of a gate operation with a parameter map.
 
         Args:
             gate_op (QuantumGate): The gate operation to transform.
-            param_map (Dict[str, Union[FloatLiteral, IntegerLiteral]]): The parameter map to use
-                                                                        for transformation.
+            param_map (dict[str, Expression]): The parameter map to use for transformation.
 
         Returns:
             None
@@ -151,13 +152,12 @@ def transform_gate_params(
             # TODO : update the arg value in expressions not just SINGLE identifiers
 
     @staticmethod
-    def get_branch_params(condition: Any) -> tuple[Optional[int], Optional[str]]:
+    def get_branch_params(condition: Any) -> tuple[int, str]:
         """
         Get the branch parameters from the branching condition
 
         Args:
-            condition (Union[UnaryExpression, BinaryExpression, IndexExpression]): The condition
-                                                                                   to analyze
+            condition (Any): The condition to analyze
 
         Returns:
             tuple[int, str]: The branch parameters
@@ -168,18 +168,36 @@ def get_branch_params(condition: Any) -> tuple[Optional[int], Optional[str]]:
                 condition.expression.collection.name,
             )
         if isinstance(condition, BinaryExpression):
-            return condition.lhs.index[0].value, condition.lhs.collection.name
+            return (
+                condition.lhs.index[0].value,
+                condition.lhs.collection.name,
+            )
         if isinstance(condition, IndexExpression):
-            return condition.index[0].value, condition.collection.name
-        return None, None
+            if isinstance(condition.index, DiscreteSet):
+                raise_qasm3_error(
+                    message="DiscreteSet not supported in branching condition",
+                    span=condition.span,
+                )
+            if isinstance(condition.index, list):
+                if isinstance(condition.index[0], RangeDefinition):
+                    raise_qasm3_error(
+                        message="RangeDefinition not supported in branching condition",
+                        span=condition.span,
+                    )
+                return (
+                    condition.index[0].value,
+                    condition.collection.name,
+                )
+        # default case
+        return -1, ""
 
     @classmethod
     def transform_function_qubits(
         cls,
         q_op: Union[QuantumGate, QuantumBarrier, QuantumReset],
-        formal_qreg_sizes: dict[str:int],
-        qubit_map: dict[tuple:tuple],
-    ) -> list:
+        formal_qreg_sizes: dict[str, int],
+        qubit_map: dict[tuple, tuple],
+    ) -> list[IndexedIdentifier]:
         """Transform the qubits of a function call to the actual qubits.
 
         Args:
@@ -191,7 +209,7 @@ def transform_function_qubits(
         Returns:
             None
         """
-        expanded_op_qubits = cls.visitor_obj._get_op_qubits(q_op, formal_qreg_sizes, qir_form=False)
+        expanded_op_qubits = cls.visitor_obj._get_op_bits(q_op, formal_qreg_sizes, qir_form=False)
 
         transformed_qubits = []
         for qubit in expanded_op_qubits:
diff --git a/qbraid_qir/qasm3/validator.py b/qbraid_qir/qasm3/validator.py
index b76d60b..69a0a23 100644
--- a/qbraid_qir/qasm3/validator.py
+++ b/qbraid_qir/qasm3/validator.py
@@ -12,17 +12,11 @@
 Module with utility functions for QASM3 visitor
 
 """
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from openqasm3.ast import ArrayType, ClassicalDeclaration
 from openqasm3.ast import IntType as Qasm3IntType
-from openqasm3.ast import (
-    QuantumGate,
-    QuantumGateDefinition,
-    ReturnStatement,
-    Statement,
-    SubroutineDefinition,
-)
+from openqasm3.ast import QuantumGate, QuantumGateDefinition, ReturnStatement, SubroutineDefinition
 
 from .elements import Variable
 from .exceptions import Qasm3ConversionError, raise_qasm3_error
@@ -53,14 +47,12 @@ def validate_register_index(index: Optional[int], size: int, qubit: bool = False
         )
 
     @staticmethod
-    def validate_statement_type(
-        blacklisted_stmts: set, statement: Statement, construct: str
-    ) -> None:
+    def validate_statement_type(blacklisted_stmts: set, statement: Any, construct: str) -> None:
         """Validate the type of a statement.
 
         Args:
             blacklisted_stmts (set): The set of blacklisted statements.
-            statement (Statement): The statement to validate.
+            statement (Any): The statement to validate.
             construct (str): The construct the statement is in.
 
         Raises:
@@ -126,10 +118,11 @@ def validate_variable_assignment_value(variable: Variable, value) -> Any:
         # For each type we will have a "castable" type set and its corresponding cast operation
         type_casted_value = qasm_variable_type_cast(qasm_type, variable.name, base_size, value)
 
+        left: Union[int, float] = 0
+        right: Union[int, float] = 0
         # check 2 - range match , if bits mentioned in base size
         if type_to_match == int:
             base_size = variable.base_size
-            left, right = 0, 0
             if qasm_type == Qasm3IntType:
                 left, right = (
                     -1 * (2 ** (base_size - 1)),
@@ -146,12 +139,11 @@ def validate_variable_assignment_value(variable: Variable, value) -> Any:
 
         elif type_to_match == float:
             base_size = variable.base_size
-            left, right = 0, 0
 
             if base_size == 32:
-                left, right = -(LIMITS_MAP["float_32"]), (LIMITS_MAP["float_32"])
+                left, right = -1.0 * (LIMITS_MAP["float_32"]), (LIMITS_MAP["float_32"])
             else:
-                left, right = -(LIMITS_MAP["float_64"]), (LIMITS_MAP["float_64"])
+                left, right = -1.0 * (LIMITS_MAP["float_64"]), (LIMITS_MAP["float_64"])
 
             if type_casted_value < left or type_casted_value > right:
                 raise_qasm3_error(
diff --git a/qbraid_qir/qasm3/visitor.py b/qbraid_qir/qasm3/visitor.py
index 03e94ae..8e23433 100644
--- a/qbraid_qir/qasm3/visitor.py
+++ b/qbraid_qir/qasm3/visitor.py
@@ -16,9 +16,9 @@
 import logging
 from abc import ABCMeta, abstractmethod
 
-# pylint: disable=too-many-instance-attributes,too-many-lines
+# pylint: disable=too-many-instance-attributes,too-many-lines,too-many-branches
 from collections import deque
-from typing import Union
+from typing import Any, Optional, Union
 
 import openqasm3.ast as qasm3_ast
 import pyqir
@@ -63,24 +63,23 @@ class BasicQasmVisitor(ProgramElementVisitor):
     """
 
     def __init__(self, initialize_runtime: bool = True, record_output: bool = True):
-        self._module = None
-        self._builder = None
-        self._entry_point = None
-        self._scope = deque([{}])
-        self._context = deque([Context.GLOBAL])
-        self._qubit_labels = {}
-        self._clbit_labels = {}
-        self._global_qreg_size_map = {}
-        self._function_qreg_size_map = deque([])  # for nested functions
-        self._function_qreg_transform_map = deque([])  # for nested functions
-        self._global_creg_size_map = {}
-        self._custom_gates = {}
-        self._subroutine_defns = {}
-        self._measured_qubits = {}
-        self._initialize_runtime = initialize_runtime
-        self._record_output = record_output
-        self._curr_scope = 0
-        self._label_scope_level = {self._curr_scope: set()}
+        self._module: pyqir.Module
+        self._builder: pyqir.Builder
+        self._entry_point: str = ""
+        self._scope: deque = deque([{}])
+        self._context: deque = deque([Context.GLOBAL])
+        self._qubit_labels: dict[str, int] = {}
+        self._clbit_labels: dict[str, int] = {}
+        self._global_qreg_size_map: dict[str, int] = {}
+        self._function_qreg_size_map: deque = deque([])  # for nested functions
+        self._function_qreg_transform_map: deque = deque([])  # for nested functions
+        self._global_creg_size_map: dict[str, int] = {}
+        self._custom_gates: dict[str, qasm3_ast.QuantumGateDefinition] = {}
+        self._subroutine_defns: dict[str, qasm3_ast.SubroutineDefinition] = {}
+        self._initialize_runtime: bool = initialize_runtime
+        self._record_output: bool = record_output
+        self._curr_scope: int = 0
+        self._label_scope_level: dict[int, set] = {self._curr_scope: set()}
 
         self._init_utilities()
 
@@ -313,10 +312,20 @@ def visit_register(
 
         current_size = len(self._qubit_labels) if is_qubit else len(self._clbit_labels)
         if is_qubit:
-            register_size = 1 if register.size is None else register.size.value
+            register_size = (
+                1 if register.size is None else register.size.value  # type: ignore[union-attr]
+            )
         else:
-            register_size = 1 if register.type.size is None else register.type.size.value
-        register_name = register.qubit.name if is_qubit else register.identifier.name
+            register_size = (
+                1
+                if register.type.size is None  # type: ignore[union-attr]
+                else register.type.size.value  # type: ignore[union-attr]
+            )
+        register_name = (
+            register.qubit.name  # type: ignore[union-attr]
+            if is_qubit
+            else register.identifier.name  # type: ignore[union-attr]
+        )
 
         size_map = self._global_qreg_size_map if is_qubit else self._global_creg_size_map
         label_map = self._qubit_labels if is_qubit else self._clbit_labels
@@ -347,7 +356,7 @@ def visit_register(
 
         logger.debug("Added labels for register '%s'", str(register))
 
-    def _check_if_name_in_scope(self, name: str, operation: any) -> None:
+    def _check_if_name_in_scope(self, name: str, operation: Any) -> None:
         """Check if a name is in scope to avoid duplicate declarations.
         Args:
             name (str): The name to check.
@@ -363,77 +372,102 @@ def _check_if_name_in_scope(self, name: str, operation: any) -> None:
             f"Variable {name} not in scope for operation {operation}", span=operation.span
         )
 
-    def _get_op_qubits(
-        self, operation: any, qreg_size_map: dict, qir_form: bool = True
-    ) -> list[Union[pyqir.qubit, qasm3_ast.IndexedIdentifier]]:
-        """Get the qubits for the operation.
+    def _get_op_bits(
+        self, operation: Any, reg_size_map: dict, qubits: bool = True, qir_form: bool = True
+    ) -> Union[list[pyqir.Constant], list[qasm3_ast.IndexedIdentifier]]:
+        """Get the quantum / classical bits for the operation.
 
         Args:
             operation (Any): The operation to get qubits for.
-            qreg_size_map (dict): The size map of the registers in scope.
-            qir_form (bool): Whether to return qubits in QIR form or not. Defaults to True.
+            reg_size_map (dict): The size map of the registers in scope.
+            qubits (bool): Whether the bits are quantum bits or classical bits. Defaults to True.
+            qir_form (bool): Whether to return bits in QIR form or not. Defaults to True.
 
         Returns:
-            list[Union[pyqir.qubit, qasm3_ast.IndexedIdentifier]]: The qubits for the operation.
+            Union[list[pyqir.Constant], list[qasm3_ast.IndexedIdentifier]] : The bits for the
+                                                                             operation.
         """
-        qir_qubits = []
-        openqasm_qubits = []
-        visited_qubits = set()
-        qubit_list = operation.qubits if isinstance(operation.qubits, list) else [operation.qubits]
-
-        for qubit in qubit_list:
-            if isinstance(qubit, qasm3_ast.IndexedIdentifier):
-                qreg_name = qubit.name.name
+        qir_bits = []
+        openqasm_bits = []
+        visited_bits = set()
+        bit_list = []
+        if isinstance(operation, qasm3_ast.QuantumMeasurementStatement):
+            assert operation.target is not None
+            bit_list = [operation.measure.qubit] if qubits else [operation.target]
+        else:
+            bit_list = (
+                operation.qubits if isinstance(operation.qubits, list) else [operation.qubits]
+            )
+
+        for bit in bit_list:
+            if isinstance(bit, qasm3_ast.IndexedIdentifier):
+                reg_name = bit.name.name
             else:
-                qreg_name = qubit.name
+                reg_name = bit.name
 
-            if qreg_name not in qreg_size_map:
+            if reg_name not in reg_size_map:
                 raise_qasm3_error(
-                    f"Missing register declaration for {qreg_name} in operation {operation}",
+                    f"Missing register declaration for {reg_name} in operation {operation}",
                     span=operation.span,
                 )
-            self._check_if_name_in_scope(qreg_name, operation)
-            qreg_size = qreg_size_map[qreg_name]
-
-            if isinstance(qubit, qasm3_ast.IndexedIdentifier):
-                if isinstance(qubit.indices[0][0], qasm3_ast.RangeDefinition):
-                    qids = Qasm3Transformer.get_qubits_from_range_definition(
-                        qubit.indices[0][0], qreg_size, is_qubit_reg=True
+            self._check_if_name_in_scope(reg_name, operation)
+
+            if isinstance(bit, qasm3_ast.IndexedIdentifier):
+                if isinstance(bit.indices[0], qasm3_ast.DiscreteSet):
+                    bit_ids = Qasm3Transformer.extract_values_from_discrete_set(bit.indices[0])
+                elif isinstance(bit.indices[0][0], qasm3_ast.RangeDefinition):
+                    bit_ids = Qasm3Transformer.get_qubits_from_range_definition(
+                        bit.indices[0][0], reg_size_map[reg_name], is_qubit_reg=qubits
                     )
                 else:
-                    qid = Qasm3ExprEvaluator.evaluate_expression(qubit.indices[0][0])
-                    Qasm3Validator.validate_register_index(qid, qreg_size, qubit=True)
-                    qids = [qid]
-                openqasm_qubits.extend(
+                    bit_id = Qasm3ExprEvaluator.evaluate_expression(bit.indices[0][0])
+                    Qasm3Validator.validate_register_index(
+                        bit_id, reg_size_map[reg_name], qubit=qubits
+                    )
+                    bit_ids = [bit_id]
+                openqasm_bits.extend(
                     [
                         qasm3_ast.IndexedIdentifier(
-                            qasm3_ast.Identifier(qreg_name), [[qasm3_ast.IntegerLiteral(i)]]
+                            qasm3_ast.Identifier(reg_name), [[qasm3_ast.IntegerLiteral(bit_id)]]
                         )
-                        for i in qids
+                        for bit_id in bit_ids
                     ]
                 )
+
             else:
-                qids = list(range(qreg_size))
-                openqasm_qubits.extend(
+                bit_ids = list(range(reg_size_map[reg_name]))
+                openqasm_bits.extend(
                     [
                         qasm3_ast.IndexedIdentifier(
-                            qasm3_ast.Identifier(qreg_name), [[qasm3_ast.IntegerLiteral(i)]]
+                            qasm3_ast.Identifier(reg_name), [[qasm3_ast.IntegerLiteral(bit_id)]]
                         )
-                        for i in qids
+                        for bit_id in bit_ids
                     ]
                 )
 
             if qir_form:
-                qreg_qids = [self._qubit_labels[f"{qreg_name}_{i}"] for i in qids]
-                for qid in qreg_qids:
-                    if qid in visited_qubits:
+                label_map = self._qubit_labels if qubits else self._clbit_labels
+                reg_ids = [label_map[f"{reg_name}_{bit_id}"] for bit_id in bit_ids]
+                for bit_id in reg_ids:
+                    if bit_id in visited_bits:
                         raise_qasm3_error(
-                            f"Duplicate qubit {qreg_name}[{qid}] argument", span=operation.span
+                            f"Duplicate {'qubit' if qubits else 'clbit'} "
+                            f"{reg_name}[{bit_id}] argument",
+                            span=operation.span,
+                        )
+                    visited_bits.add(bit_id)
+                qir_bits.extend(
+                    [
+                        (
+                            pyqir.qubit(self._module.context, bit_id)
+                            if qubits
+                            else pyqir.result(self._module.context, bit_id)
                         )
-                    visited_qubits.add(qid)
-                qir_qubits.extend([pyqir.qubit(self._module.context, n) for n in qreg_qids])
+                        for bit_id in reg_ids
+                    ]
+                )
 
-        return qir_qubits if qir_form else openqasm_qubits
+        return qir_bits if qir_form else openqasm_bits
 
     def _visit_measurement(self, statement: qasm3_ast.QuantumMeasurementStatement) -> None:
         """Visit a measurement statement element.
@@ -445,36 +479,25 @@ def _visit_measurement(self, statement: qasm3_ast.QuantumMeasurementStatement) -
             None
         """
         logger.debug("Visiting measurement statement '%s'", str(statement))
+
         source = statement.measure.qubit
         target = statement.target
-        source_id, target_id = None, None
-        # TODO: handle in-function measurements
-        source_name = source.name
-        if isinstance(source, qasm3_ast.IndexedIdentifier):
-            source_name = source.name.name
-            if isinstance(source.indices[0][0], qasm3_ast.RangeDefinition):
-                raise_qasm3_error(
-                    f"Range based measurement {statement} not supported at the moment",
-                    span=statement.span,
-                )
-            source_id = source.indices[0][0].value
-
-        target_name = target.name
-        if isinstance(target, qasm3_ast.IndexedIdentifier):
-            target_name = target.name.name
-            if isinstance(target.indices[0][0], qasm3_ast.RangeDefinition):
-                raise_qasm3_error(
-                    f"Range based measurement {statement} not supported at the moment",
-                    span=statement.span,
-                )
-            target_id = target.indices[0][0].value
+        assert source and target
 
+        # # TODO: handle in-function measurements
+        source_name: str = (
+            source.name if isinstance(source, qasm3_ast.Identifier) else source.name.name
+        )
         if source_name not in self._global_qreg_size_map:
             raise_qasm3_error(
                 f"Missing register declaration for {source_name} in measurement "
                 f"operation {statement}",
                 span=statement.span,
             )
+
+        target_name: str = (
+            target.name if isinstance(target, qasm3_ast.Identifier) else target.name.name
+        )
         if target_name not in self._global_creg_size_map:
             raise_qasm3_error(
                 f"Missing register declaration for {target_name} in measurement "
@@ -482,41 +505,22 @@ def _visit_measurement(self, statement: qasm3_ast.QuantumMeasurementStatement) -
                 span=statement.span,
             )
 
-        def _build_qir_measurement(
-            src_name: str,
-            src_id: Union[int, None],
-            target_name: str,
-            target_id: Union[int, None],
-        ):
-            src_id = 0 if src_id is None else src_id
-            target_id = 0 if target_id is None else target_id
+        source_ids = self._get_op_bits(
+            statement, reg_size_map=self._global_qreg_size_map, qubits=True
+        )
+        target_ids = self._get_op_bits(
+            statement, reg_size_map=self._global_creg_size_map, qubits=False
+        )
 
-            source_qubit = pyqir.qubit(
-                self._module.context, self._qubit_labels[f"{src_name}_{src_id}"]
-            )
-            result = pyqir.result(
-                self._module.context,
-                self._clbit_labels[f"{target_name}_{target_id}"],
+        if len(source_ids) != len(target_ids):
+            raise_qasm3_error(
+                f"Register sizes of {source_name} and {target_name} do not match "
+                "for measurement operation",
+                span=statement.span,
             )
-            pyqir._native.mz(self._builder, source_qubit, result)
 
-        if source_id is None and target_id is None:
-            if self._global_qreg_size_map[source_name] != self._global_creg_size_map[target_name]:
-                raise_qasm3_error(
-                    f"Register sizes of {source_name} and {target_name} do not match "
-                    "for measurement operation",
-                    span=statement.span,
-                )
-            for i in range(self._global_qreg_size_map[source_name]):
-                _build_qir_measurement(source_name, i, target_name, i)
-        else:
-            Qasm3Validator.validate_register_index(
-                source_id, self._global_qreg_size_map[source_name], qubit=True
-            )
-            Qasm3Validator.validate_register_index(
-                target_id, self._global_creg_size_map[target_name], qubit=False
-            )
-            _build_qir_measurement(source_name, source_id, target_name, target_id)
+        for src_id, tgt_id in zip(source_ids, target_ids):
+            pyqir._native.mz(self._builder, src_id, tgt_id)  # type: ignore[arg-type]
 
     def _visit_reset(self, statement: qasm3_ast.QuantumReset) -> None:
         """Visit a reset statement element.
@@ -530,15 +534,18 @@ def _visit_reset(self, statement: qasm3_ast.QuantumReset) -> None:
         logger.debug("Visiting reset statement '%s'", str(statement))
         if len(self._function_qreg_size_map) > 0:  # atleast in SOME function scope
             # transform qubits to use the global qreg identifiers
-            statement.qubits = Qasm3Transformer.transform_function_qubits(
-                statement,
-                self._function_qreg_size_map[-1],
-                self._function_qreg_transform_map[-1],
+            statement.qubits = (
+                Qasm3Transformer.transform_function_qubits(  # type: ignore[assignment]
+                    statement,
+                    self._function_qreg_size_map[-1],
+                    self._function_qreg_transform_map[-1],
+                )
             )
-        qubit_ids = self._get_op_qubits(statement, self._global_qreg_size_map, True)
+        qubit_ids = self._get_op_bits(statement, self._global_qreg_size_map, True)
 
         for qid in qubit_ids:
-            pyqir._native.reset(self._builder, qid)
+            # qid is of type Constant which is inherited from Value, so we ignore the type error
+            pyqir._native.reset(self._builder, qid)  # type: ignore[arg-type]
 
     def _visit_barrier(self, barrier: qasm3_ast.QuantumBarrier) -> None:
         """Visit a barrier statement element.
@@ -552,13 +559,17 @@ def _visit_barrier(self, barrier: qasm3_ast.QuantumBarrier) -> None:
         # if barrier is applied to ALL qubits at once, we are fine
         if len(self._function_qreg_size_map) > 0:  # atleast in SOME function scope
             # transform qubits to use the global qreg identifiers
-            barrier.qubits = Qasm3Transformer.transform_function_qubits(
-                barrier,
-                self._function_qreg_size_map[-1],
-                self._function_qreg_transform_map[-1],
-            )
 
-        barrier_qubits = self._get_op_qubits(barrier, self._global_qreg_size_map)
+            # since we are changing the qubits to IndexedIdentifiers, we need to supress the
+            # error for the type checker
+            barrier.qubits = (
+                Qasm3Transformer.transform_function_qubits(  # type: ignore [assignment]
+                    barrier,
+                    self._function_qreg_size_map[-1],
+                    self._function_qreg_transform_map[-1],
+                )
+            )
+        barrier_qubits = self._get_op_bits(barrier, self._global_qreg_size_map)
         total_qubit_count = sum(self._global_qreg_size_map.values())
         if len(barrier_qubits) == total_qubit_count:
             pyqir._native.barrier(self._builder)
@@ -627,7 +638,7 @@ def _visit_basic_gate_operation(
 
         logger.debug("Visiting basic gate operation '%s'", str(operation))
         op_name: str = operation.name.name
-        op_qubits = self._get_op_qubits(operation, self._global_qreg_size_map)
+        op_qubits = self._get_op_bits(operation, self._global_qreg_size_map)
         inverse_action = None
         if not inverse:
             qir_func, op_qubit_count = map_qasm_op_to_pyqir_callable(op_name)
@@ -676,7 +687,11 @@ def _visit_custom_gate_operation(
         logger.debug("Visiting custom gate operation '%s'", str(operation))
         gate_name: str = operation.name.name
         gate_definition: qasm3_ast.QuantumGateDefinition = self._custom_gates[gate_name]
-        op_qubits = self._get_op_qubits(operation, self._global_qreg_size_map, qir_form=False)
+        op_qubits: list[qasm3_ast.IndexedIdentifier] = (
+            self._get_op_bits(  # type: ignore [assignment]
+                operation, self._global_qreg_size_map, qir_form=False
+            )
+        )
 
         Qasm3Validator.validate_gate_call(operation, gate_definition, len(op_qubits))
         # we need this because the gates applied inside a gate definition use the
@@ -700,15 +715,14 @@ def _visit_custom_gate_operation(
         self._push_context(Context.GATE)
 
         for gate_op in gate_definition_ops:
-            if gate_op.name.name == gate_name:
-                raise_qasm3_error(
-                    f"Recursive definitions not allowed for gate {gate_name}", span=gate_op.span
-                )
-
-            # necessary to avoid modifying the original gate definition
-            # in case the gate is reapplied
-            gate_op_copy = copy.deepcopy(gate_op)
             if isinstance(gate_op, qasm3_ast.QuantumGate):
+                # necessary to avoid modifying the original gate definition
+                # in case the gate is reapplied
+                gate_op_copy = copy.deepcopy(gate_op)
+                if gate_op.name.name == gate_name:
+                    raise_qasm3_error(
+                        f"Recursive definitions not allowed for gate {gate_name}", span=gate_op.span
+                    )
                 Qasm3Transformer.transform_gate_params(gate_op_copy, param_map)
                 Qasm3Transformer.transform_gate_qubits(gate_op_copy, qubit_map)
                 # need to trickle the inverse down to the child gates
@@ -776,10 +790,12 @@ def _visit_generic_gate_operation(self, operation: qasm3_ast.QuantumGate) -> Non
         if not self._in_gate_scope() and len(self._function_qreg_size_map) > 0:
             # we are in SOME function scope
             # transform qubits to use the global qreg identifiers
-            operation.qubits = Qasm3Transformer.transform_function_qubits(
-                operation,
-                self._function_qreg_size_map[-1],
-                self._function_qreg_transform_map[-1],
+            operation.qubits = (
+                Qasm3Transformer.transform_function_qubits(  # type: ignore [assignment]
+                    operation,
+                    self._function_qreg_size_map[-1],
+                    self._function_qreg_transform_map[-1],
+                )
             )
         # Applying the inverse first and then the power is same as
         # apply the power first and then inverting the result
@@ -816,14 +832,16 @@ def _visit_constant_declaration(self, statement: qasm3_ast.ConstantDeclaration)
         base_type = statement.type
         if isinstance(base_type, qasm3_ast.BoolType):
             base_size = 1
-        elif base_type.size is None:
-            base_size = 32  # default for now
-        else:
-            base_size = Qasm3ExprEvaluator.evaluate_expression(base_type.size, const_expr=True)
-            if not isinstance(base_size, int) or base_size <= 0:
-                raise_qasm3_error(
-                    f"Invalid base size {base_size} for variable {var_name}", span=statement.span
-                )
+        elif hasattr(base_type, "size"):
+            if base_type.size is None:
+                base_size = 32  # default for now
+            else:
+                base_size = Qasm3ExprEvaluator.evaluate_expression(base_type.size, const_expr=True)
+                if not isinstance(base_size, int) or base_size <= 0:
+                    raise_qasm3_error(
+                        f"Invalid base size {base_size} for variable {var_name}",
+                        span=statement.span,
+                    )
 
         variable = Variable(var_name, base_type, base_size, [], init_value, is_constant=True)
 
@@ -887,6 +905,7 @@ def _visit_classical_declaration(self, statement: qasm3_ast.ClassicalDeclaration
 
             init_value = None
             for dim in reversed(final_dimensions):
+                assert isinstance(dim, int)
                 init_value = [init_value for _ in range(dim)]
 
         if statement.init_expression:
@@ -900,7 +919,7 @@ def _visit_classical_declaration(self, statement: qasm3_ast.ClassicalDeclaration
         if not isinstance(base_type, qasm3_ast.BoolType):
             base_size = (
                 32
-                if base_type.size is None
+                if not hasattr(base_type, "size") or base_type.size is None
                 else Qasm3ExprEvaluator.evaluate_expression(base_type.size)
             )
 
@@ -918,6 +937,7 @@ def _visit_classical_declaration(self, statement: qasm3_ast.ClassicalDeclaration
 
         if statement.init_expression:
             if isinstance(init_value, list):
+                assert variable.dims is not None
                 Qasm3Validator.validate_array_assignment_values(variable, variable.dims, init_value)
             else:
                 variable.value = Qasm3Validator.validate_variable_assignment_value(
@@ -938,46 +958,48 @@ def _visit_classical_assignment(self, statement: qasm3_ast.ClassicalAssignment)
         lvalue = statement.lvalue
         var_name = lvalue.name
 
-        if isinstance(lvalue, qasm3_ast.IndexedIdentifier):
+        if isinstance(var_name, qasm3_ast.Identifier):
             var_name = var_name.name
 
         var = self._get_from_visible_scope(var_name)
 
-        if var is None:
+        if var is None:  # we check for none here, so type errors are irrelevant afterwards
             raise_qasm3_error(f"Undefined variable {var_name} in assignment", span=statement.span)
 
-        if var.is_constant:
+        if var.is_constant:  # type: ignore[union-attr]
             raise_qasm3_error(
                 f"Assignment to constant variable {var_name} not allowed", span=statement.span
             )
 
         var_value = Qasm3ExprEvaluator.evaluate_expression(statement.rvalue)
 
-        # currently we support single array assignment only
-        # range based assignment not supported yet
+        # currently we support single array assignment only.
+        # TODO: range based assignment
 
         # cast + validation
-        var_value = Qasm3Validator.validate_variable_assignment_value(var, var_value)
-
+        var_value = Qasm3Validator.validate_variable_assignment_value(
+            var, var_value  # type: ignore[arg-type]
+        )
         # handle assignment for arrays
         if isinstance(lvalue, qasm3_ast.IndexedIdentifier):
             # stupid indices structure in openqasm :/
-            if len(lvalue.indices[0]) > 1:
+            if len(lvalue.indices[0]) > 1:  # type: ignore[arg-type]
                 indices = lvalue.indices[0]
             else:
-                indices = [idx[0] for idx in lvalue.indices]
+                indices = [idx[0] for idx in lvalue.indices]  # type: ignore[assignment, index]
 
             validated_indices = Qasm3Analyzer.analyze_classical_indices(
-                indices, self._get_from_visible_scope(var_name)
+                indices, self._get_from_visible_scope(var_name)  # type: ignore[arg-type]
+            )
+            Qasm3Transformer.update_array_element(
+                var.value, validated_indices, var_value  # type: ignore[union-attr, arg-type]
             )
-            Qasm3Transformer.update_array_element(var.value, validated_indices, var_value)
         else:
-            var.value = var_value
-
-        self._update_var_in_scope(var)
+            var.value = var_value  # type: ignore[union-attr]
+        self._update_var_in_scope(var)  # type: ignore[arg-type]
 
     def _evaluate_array_initialization(
-        self, array_literal: qasm3_ast.ArrayLiteral, dimensions: list[int], base_type: any
+        self, array_literal: qasm3_ast.ArrayLiteral, dimensions: list[int], base_type: Any
     ) -> list:
         """Evaluate an array initialization.
 
@@ -1078,7 +1100,7 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> None:
                 f"Unexpected type {type(statement.set_declaration)} of set_declaration in loop."
             )
 
-        i = None  # will store iteration Variable to update to loop scope
+        i: Optional[Variable]  # will store iteration Variable to update to loop scope
 
         for ival in irange:
             self._push_context(Context.BLOCK)
@@ -1093,8 +1115,9 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> None:
             i = self._get_from_visible_scope(statement.identifier.name)
 
             # Update scope with current value of loop Variable
-            i.value = ival
-            self._update_var_in_scope(i)
+            if i is not None:
+                i.value = ival
+                self._update_var_in_scope(i)
 
             for stmt in statement.block:
                 self.visit_statement(stmt)
@@ -1198,9 +1221,9 @@ def _visit_function_call(self, statement: qasm3_ast.FunctionCall) -> None:
                 span=statement.span,
             )
 
-        duplicate_qubit_detect_map = {}
-        qubit_transform_map = {}  # {(formal arg, idx) : (actual arg, idx)}
-        formal_qreg_size_map = {}
+        duplicate_qubit_detect_map: dict = {}
+        qubit_transform_map: dict = {}  # {(formal arg, idx) : (actual arg, idx)}
+        formal_qreg_size_map: dict = {}
 
         quantum_vars, classical_vars = [], []
 
@@ -1329,7 +1352,9 @@ def _process_quantum_arg(formal_arg, actual_arg, formal_reg_name, actual_arg_nam
             actual_arg_name = None
             if isinstance(actual_arg, qasm3_ast.Identifier):
                 actual_arg_name = actual_arg.name
-            elif isinstance(actual_arg, qasm3_ast.IndexExpression):
+            elif isinstance(actual_arg, qasm3_ast.IndexExpression) and isinstance(
+                actual_arg.collection, qasm3_ast.Identifier
+            ):
                 actual_arg_name = actual_arg.collection.name
 
             if isinstance(formal_arg, qasm3_ast.ClassicalArgument):
@@ -1391,10 +1416,10 @@ def _visit_alias_statement(self, statement: qasm3_ast.AliasStatement) -> None:
         target = statement.target
         value = statement.value
 
-        alias_reg_name = target.name
-        alias_reg_size = None
-        aliased_reg_name = None
-        aliased_reg_size = None
+        alias_reg_name: str = target.name
+        alias_reg_size: int = 0
+        aliased_reg_name: str = ""
+        aliased_reg_size: int = 0
 
         # Alias should not be redeclared earlier as a variable or a constant
         if self._check_in_scope(alias_reg_name):
@@ -1403,7 +1428,9 @@ def _visit_alias_statement(self, statement: qasm3_ast.AliasStatement) -> None:
 
         if isinstance(value, qasm3_ast.Identifier):
             aliased_reg_name = value.name
-        elif isinstance(value, qasm3_ast.IndexExpression):
+        elif isinstance(value, qasm3_ast.IndexExpression) and isinstance(
+            value.collection, qasm3_ast.Identifier
+        ):
             aliased_reg_name = value.collection.name
         else:
             raise_qasm3_error(f"Unsupported aliasing {statement}", span=statement.span)
@@ -1469,11 +1496,11 @@ def _visit_switch_statement(self, statement: qasm3_ast.SwitchStatement) -> None:
         """
         # 1. analyze the target - it should ONLY be int, not casted
         switch_target = statement.target
-
+        switch_target_name = ""
         # either identifier or indexed expression
         if isinstance(switch_target, qasm3_ast.Identifier):
             switch_target_name = switch_target.name
-        else:
+        elif isinstance(switch_target, qasm3_ast.IndexExpression):
             switch_target_name, _ = Qasm3Analyzer.analyze_index_expression(switch_target)
 
         if not Qasm3Validator.validate_variable_type(
@@ -1571,7 +1598,7 @@ def visit_statement(self, statement: qasm3_ast.Statement) -> None:
         visitor_function = visit_map.get(type(statement))
 
         if visitor_function:
-            visitor_function(statement)
+            visitor_function(statement)  # type: ignore[operator]
         else:
             raise_qasm3_error(
                 f"Unsupported statement of type {type(statement)}", span=statement.span
@@ -1581,4 +1608,4 @@ def ir(self) -> str:
         return str(self._module)
 
     def bitcode(self) -> bytes:
-        return self._module.bitcode()
+        return self._module.bitcode
diff --git a/tox.ini b/tox.ini
index 400e067..07c052e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -48,6 +48,13 @@ deps = black[jupyter]
 commands =
     black qbraid_qir tools tests examples {posargs}
 
+[testenv:mypy]
+envdir = .tox/linters
+skip_install = true
+deps = mypy
+commands =
+    mypy qbraid_qir
+
 [testenv:headers]
 envdir = .tox/linters
 skip_install = true
@@ -76,9 +83,11 @@ deps =
     {[testenv:pylint]deps}
     {[testenv:isort]deps}
     {[testenv:black]deps}
+    {[testenv:mypy]deps}
     {[testenv:headers]deps}
 commands =
     {[testenv:pylint]commands}
     {[testenv:isort]commands} {posargs:--check-only}
     {[testenv:black]commands} {posargs:--check}
+    {[testenv:mypy]commands}
     {[testenv:headers]commands}
\ No newline at end of file


Test patch:
diff --git a/tests/qasm3_qir/converter/declarations/test_classical.py b/tests/qasm3_qir/converter/declarations/test_classical.py
index 1b587c7..cd33696 100644
--- a/tests/qasm3_qir/converter/declarations/test_classical.py
+++ b/tests/qasm3_qir/converter/declarations/test_classical.py
@@ -293,8 +293,8 @@ def test_array_expressions():
     arr_float64[0][1] = d*d;
 
     qubit q;
-    rx(arr_int[0][1] + arr_uint[0][1]) q;
-    rx(arr_float32[0][1] + arr_float64[0][1]) q;
+    rx(arr_int[{0,1}] + arr_uint[0][1]) q;
+    rx(arr_float32[{0,1}] + arr_float64[0][1]) q;
     """
 
     a = 2
diff --git a/tests/qasm3_qir/converter/test_if.py b/tests/qasm3_qir/converter/test_if.py
index 0969de4..7f17991 100644
--- a/tests/qasm3_qir/converter/test_if.py
+++ b/tests/qasm3_qir/converter/test_if.py
@@ -205,3 +205,42 @@ def test_incorrect_if():
            }
            """
         )
+    with pytest.raises(
+        Qasm3ConversionError,
+        match=r"RangeDefinition not supported in branching condition",
+    ):
+        _ = qasm3_to_qir(
+            """
+            OPENQASM 3;
+           include "stdgates.inc";
+           qubit[2] q;
+           bit[2] c;
+
+           h q;
+           measure q->c;
+
+           if(c[0:1]){
+            cx q;
+           }
+           """
+        )
+
+    with pytest.raises(
+        Qasm3ConversionError,
+        match=r"DiscreteSet not supported in branching condition",
+    ):
+        _ = qasm3_to_qir(
+            """
+            OPENQASM 3;
+           include "stdgates.inc";
+           qubit[2] q;
+           bit[2] c;
+
+           h q;
+           measure q->c;
+
+           if(c[{0,1}]){
+            cx q;
+           }
+           """
+        )
diff --git a/tests/qasm3_qir/converter/test_measurement.py b/tests/qasm3_qir/converter/test_measurement.py
index 8c675c6..a37a42d 100644
--- a/tests/qasm3_qir/converter/test_measurement.py
+++ b/tests/qasm3_qir/converter/test_measurement.py
@@ -34,15 +34,18 @@ def test_measure():
     c1 = measure q1;
     measure q1 -> c1;
     c2[0] = measure q3[0];
+    measure q1[:1] -> c1[1];
+    measure q2[{0, 1}] -> c1[{1, 0}];
+
     """
 
     result = qasm3_to_qir(qasm3_string)
     generated_qir = str(result).splitlines()
     check_attributes(generated_qir, 8, 3)
-    qubit_list = [0, 1, 0, 1, 7]
-    bit_list = [0, 1, 0, 1, 2]
+    qubit_list = [0, 1, 0, 1, 7, 0, 2, 3]
+    bit_list = [0, 1, 0, 1, 2, 1, 1, 0]
 
-    check_measure_op(generated_qir, 5, qubit_list, bit_list)
+    check_measure_op(generated_qir, 8, qubit_list, bit_list)
 
 
 def test_incorrect_measure():
@@ -50,27 +53,6 @@ def run_test(qasm3_code, error_message):
         with pytest.raises(Qasm3ConversionError, match=error_message):
             _ = qasm3_to_qir(qasm3_code)
 
-    # Test for range based measurement not supported
-    run_test(
-        """
-        OPENQASM 3;
-        qubit[2] q1;
-        bit[2] c1;
-        measure q1[0:2] -> c1[:2];  // not supported 
-    """,
-        r"Range based measurement .* not supported at the moment",
-    )
-
-    run_test(
-        """
-        OPENQASM 3;
-        qubit[2] q1;
-        bit[2] c1;
-        measure q1 -> c1[:2];  // not supported 
-    """,
-        r"Range based measurement .* not supported at the moment",
-    )
-
     # Test for undeclared register q2
     run_test(
         """
@@ -105,6 +87,18 @@ def run_test(qasm3_code, error_message):
         r"Register sizes of q1 and c2 do not match .*",
     )
 
+    # Test for size mismatch between q1 and c2 in ranges
+    run_test(
+        """
+        OPENQASM 3;
+        qubit[5] q1;
+        bit[4] c1;
+        bit[1] c2;
+        c1[:3] = measure q1;  // size mismatch
+    """,
+        r"Register sizes of q1 and c1 do not match .*",
+    )
+
     # Test for out of bounds index for q1
     run_test(
         """
diff --git a/tests/qir_utils.py b/tests/qir_utils.py
index 9ece70a..1690e0f 100644
--- a/tests/qir_utils.py
+++ b/tests/qir_utils.py
@@ -14,7 +14,7 @@
 """
 
 import struct
-from typing import List, Union
+from typing import Union
 
 from pyqir import (
     Context,
@@ -103,7 +103,7 @@ def reset_call_string(qb: int) -> str:
     return f"call void @__quantum__qis__reset__body({_qubit_string(qb)})"
 
 
-def generic_op_call_string(name: str, qbs: List[int]) -> str:
+def generic_op_call_string(name: str, qbs: list[int]) -> str:
     args = ", ".join(_qubit_string(qb) for qb in qbs)
     return f"call void @__quantum__qis__{name}__body({args})"
 
@@ -118,7 +118,7 @@ def get_entry_point(mod: Module) -> Function:
     return func
 
 
-def get_entry_point_body(qir: List[str]) -> List[str]:
+def get_entry_point_body(qir: list[str]) -> list[str]:
     joined = "\n".join(qir)
     mod = Module.from_ir(Context(), joined)
     func = next(filter(is_entry_point, mod.functions))
@@ -144,7 +144,7 @@ def check_attributes_on_entrypoint(
     ), f"Incorrect result count: {expected_results} expected, {actual_results} actual"
 
 
-def check_attributes(qir: List[str], expected_qubits: int = 0, expected_results: int = 0) -> None:
+def check_attributes(qir: list[str], expected_qubits: int = 0, expected_results: int = 0) -> None:
     x = "\n".join(qir)
     mod = Module.from_ir(Context(), x)
     func = next(filter(is_entry_point, mod.functions))
@@ -152,7 +152,7 @@ def check_attributes(qir: List[str], expected_qubits: int = 0, expected_results:
     check_attributes_on_entrypoint(func, expected_qubits, expected_results)
 
 
-def check_resets(qir: List[str], expected_resets: int, qubit_list: List[int]):
+def check_resets(qir: list[str], expected_resets: int, qubit_list: list[int]):
     entry_body = get_entry_point_body(qir)
     reset_count = 0
     for line in entry_body:
@@ -169,7 +169,7 @@ def check_resets(qir: List[str], expected_resets: int, qubit_list: List[int]):
         assert False, f"Incorrect reset count: {expected_resets} expected, {reset_count} actual"
 
 
-def check_barrier(qir: List[str], expected_barriers: int):
+def check_barrier(qir: list[str], expected_barriers: int):
     entry_body = get_entry_point_body(qir)
     barrier_count = 0
     for line in entry_body:
@@ -185,7 +185,7 @@ def check_barrier(qir: List[str], expected_barriers: int):
         ), f"Incorrect barrier count: {expected_barriers} expected, {barrier_count} actual"
 
 
-def check_measure_op(qir: List[str], expected_ops: int, qubit_list: List[int], bit_list: List[int]):
+def check_measure_op(qir: list[str], expected_ops: int, qubit_list: list[int], bit_list: list[int]):
     entry_body = get_entry_point_body(qir)
     measure_count = 0
     q_id, b_id = 0, 0
@@ -209,7 +209,7 @@ def check_measure_op(qir: List[str], expected_ops: int, qubit_list: List[int], b
 
 
 def check_single_qubit_gate_op(
-    qir: List[str], expected_ops: int, qubit_list: List[int], gate_name: str
+    qir: list[str], expected_ops: int, qubit_list: list[int], gate_name: str
 ):
     entry_body = get_entry_point_body(qir)
     op_count = 0
@@ -237,7 +237,7 @@ def check_single_qubit_gate_op(
 
 
 def check_two_qubit_gate_op(
-    qir: List[str], expected_ops: int, qubit_lists: List[int], gate_name: str
+    qir: list[str], expected_ops: int, qubit_lists: list[int], gate_name: str
 ):
     entry_body = get_entry_point_body(qir)
     op_count = 0
@@ -263,7 +263,7 @@ def check_two_qubit_gate_op(
 
 # pylint: disable-next=too-many-locals
 def check_single_qubit_u3_op(
-    entry_body: List[str], expected_ops: int, qubit_list: List[int], param_list: List[float]
+    entry_body: list[str], expected_ops: int, qubit_list: list[int], param_list: list[float]
 ):
     theta, phi, lam = param_list
     op_count = 0
@@ -306,10 +306,10 @@ def check_single_qubit_u3_op(
 
 
 def check_single_qubit_rotation_op(
-    qir: List[str],
+    qir: list[str],
     expected_ops: int,
-    qubit_list: List[int],
-    param_list: List[float],
+    qubit_list: list[int],
+    param_list: list[float],
     gate_name: str,
 ):
     entry_body = get_entry_point_body(qir)
@@ -338,7 +338,7 @@ def check_single_qubit_rotation_op(
 
 
 def check_three_qubit_gate_op(
-    qir: List[str], expected_ops: int, qubit_lists: List[int], gate_name: str
+    qir: list[str], expected_ops: int, qubit_lists: list[int], gate_name: str
 ):
     entry_body = get_entry_point_body(qir)
     op_count = 0
@@ -361,7 +361,7 @@ def check_three_qubit_gate_op(
         ), f"Incorrect three qubit gate count: {expected_ops} expected, {op_count} actual"
 
 
-def _validate_simple_custom_op(entry_body: List[str]):
+def _validate_simple_custom_op(entry_body: list[str]):
     custom_op_lines = [
         initialize_call_string(),
         single_op_call_string("h", 0),
@@ -378,7 +378,7 @@ def _validate_simple_custom_op(entry_body: List[str]):
         assert body_line.strip() == custom_op_lines[i].strip(), "Incorrect custom op line"
 
 
-def _validate_nested_custom_op(entry_body: List[str]):
+def _validate_nested_custom_op(entry_body: list[str]):
     nested_op_lines = [
         initialize_call_string(),
         single_op_call_string("h", 1),
@@ -397,7 +397,7 @@ def _validate_nested_custom_op(entry_body: List[str]):
         assert body_line.strip() == nested_op_lines[i].strip(), "Incorrect nested op line"
 
 
-def _validate_complex_custom_op(entry_body: List[str]):
+def _validate_complex_custom_op(entry_body: list[str]):
     complex_op_lines = [
         initialize_call_string(),
         single_op_call_string("h", 0),
@@ -416,7 +416,7 @@ def _validate_complex_custom_op(entry_body: List[str]):
         assert body_line.strip() == complex_op_lines[i].strip(), "Incorrect complex op line"
 
 
-def check_custom_qasm_gate_op(qir: List[str], test_type: str):
+def check_custom_qasm_gate_op(qir: list[str], test_type: str):
     entry_body = get_entry_point_body(qir)
     print(entry_body)
     if test_type == "simple":
@@ -430,7 +430,7 @@ def check_custom_qasm_gate_op(qir: List[str], test_type: str):
 
 
 def check_expressions(
-    qir: List[str], expected_ops: int, gates: List[str], expression_values, qubits: List[int]
+    qir: list[str], expected_ops: int, gates: list[str], expression_values, qubits: list[int]
 ):
     entry_body = get_entry_point_body(qir)
     op_count = 0
@@ -452,12 +452,12 @@ def check_expressions(
 
 
 def check_simple_if(
-    qir: List[str],  # pylint: disable=unused-argument
+    qir: list[str],  # pylint: disable=unused-argument
 ):
     pass
 
 
 def check_complex_if(
-    qir: List[str],  # pylint: disable=unused-argument
+    qir: list[str],  # pylint: disable=unused-argument
 ):
     pass