diff --git a/CHANGELOG.md b/CHANGELOG.md
index b19335f..20919f4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,7 @@ Types of changes:
 
 ### Improved / Modified
 - Consolidated the hardcoded `"__PYQASM_QUBITS__"` string literals scattered across `visitor.py`, `transformer.py` and `pulse/utils.py` into a single `INTERNAL_QUBIT_REGISTER` constant in `elements.py`, alongside an `is_internal_qubit_register()` helper that is now the one place the internal register is recognised. ([#325](https://github.com/qBraid/pyqasm/pull/325))
+- Added / updated type hinting for `base.py`, `qasm2.py`, `qasm3.py`, and `visitor.py` signatures, plus the `visit_statement` / `visit_basic_block` signatures in `pulse/visitor.py`. Fixed grammatical typos in `base.py`, `qasm2.py`, `qasm3.py`, and `visitor.py` docstrings. Fixes incorrect return types in the docstrings of `visitor.py`. ([#346](https://github.com/qBraid/pyqasm/pull/346))
 
 ### Deprecated
 
diff --git a/src/pyqasm/modules/base.py b/src/pyqasm/modules/base.py
index a6b0903..d91362d 100644
--- a/src/pyqasm/modules/base.py
+++ b/src/pyqasm/modules/base.py
@@ -22,14 +22,14 @@
 from abc import ABC, abstractmethod
 from collections import Counter
 from copy import deepcopy
-from typing import Iterator, Optional, Sequence, TypeVar
+from typing import Any, Callable, Iterator, Optional, Sequence, TypeVar
 
 import openqasm3.ast as qasm3_ast
 from openqasm3.ast import BranchingStatement, Program, QuantumGate
 
 from pyqasm.analyzer import Qasm3Analyzer
 from pyqasm.decomposer import Decomposer
-from pyqasm.elements import ClbitDepthNode, QubitDepthNode
+from pyqasm.elements import BasisSet, ClbitDepthNode, QubitDepthNode
 from pyqasm.exceptions import UnrollError, ValidationError
 from pyqasm.maps import QUANTUM_STATEMENTS
 from pyqasm.maps.decomposition_rules import DECOMPOSITION_RULES
@@ -90,7 +90,10 @@ def drop_statements(statements: list[StatementT], unwanted: type) -> list[Statem
     return kept
 
 
-def track_user_operation(func):
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def track_user_operation(func: F) -> F:
     """Decorator to track user operations on a QasmModule."""
 
     @functools.wraps(func)
@@ -113,11 +116,11 @@ def wrapper(self, *args, **kwargs):
         self._user_operations.append(log_message)
         return func(self, *args, **kwargs)
 
-    return wrapper
+    return wrapper  # type: ignore
 
 
 class QasmModule(ABC):  # pylint: disable=too-many-instance-attributes, too-many-public-methods
-    """Abstract class for a Qasm module
+    """Abstract class for a Qasm module.
 
     Args:
         name (str): Name of the module.
@@ -170,18 +173,15 @@ def num_qubits(self) -> int:
         return self._num_qubits
 
     @num_qubits.setter
-    def num_qubits(self, value: int):
-        """Setter for the number of qubits"""
+    def num_qubits(self, value: int) -> None:
+        """Setter for the number of qubits."""
         self._num_qubits = value
 
-    def _add_qubit_register(self, reg_name: str, num_qubits: int):
-        """Add qubits to the module
+    def _add_qubit_register(self, reg_name: str, num_qubits: int) -> None:
+        """Add qubits to the module.
 
         Args:
-            num_qubits (int): The number of qubits to add to the module
-
-        Returns:
-            None
+            num_qubits (int): The number of qubits to add to the module.
         """
         self._qubit_registers[reg_name] = num_qubits
         self._num_qubits += num_qubits
@@ -195,39 +195,40 @@ def num_clbits(self) -> int:
         return self._num_clbits
 
     @num_clbits.setter
-    def num_clbits(self, value: int):
-        """Setter for the number of classical bits"""
+    def num_clbits(self, value: int) -> None:
+        """Setter for the number of classical bits."""
         self._num_clbits = value
 
-    def _add_classical_register(self, reg_name: str, num_clbits: int):
-        """Add classical bits to the module
+    def _add_classical_register(self, reg_name: str, num_clbits: int) -> None:
+        """Add classical bits to the module.
 
         Args:
-            num_clbits (int): The number of classical bits to add to the module
-
-        Returns:
-            None
+            num_clbits (int): The number of classical bits to add to the module.
         """
         self._classical_registers[reg_name] = num_clbits
         self._num_clbits += num_clbits
 
     @property
     def original_program(self) -> Program:
-        """Returns the program AST for the original qasm supplied by the user"""
+        """Returns the program AST for the original qasm supplied by the user."""
         return self._original_program
 
     @property
     def unrolled_ast(self) -> Program:
-        """Returns the unrolled AST for the module"""
+        """Returns the unrolled AST for the module."""
         return self._unrolled_ast
 
     @unrolled_ast.setter
-    def unrolled_ast(self, value: Program):
-        """Setter for the unrolled AST"""
+    def unrolled_ast(self, value: Program) -> None:
+        """Setter for the unrolled AST."""
         self._unrolled_ast = value
 
     def has_measurements(self) -> bool:
-        """Check if the module has any measurement operations."""
+        """Check if the module has any measurement operations.
+
+        Returns:
+            bool: True if the module has measurement operations, False otherwise.
+        """
         if self._has_measurements is None:
             self._has_measurements = False
             # try to check in the unrolled version as that will a better indicator of
@@ -245,13 +246,13 @@ def has_measurements(self) -> bool:
 
     @track_user_operation
     def remove_measurements(self, in_place: bool = True) -> Optional["QasmModule"]:
-        """Remove the measurement operations
+        """Remove the measurement operations.
 
         Args:
             in_place (bool): Flag to indicate if the removal should be done in place.
 
         Returns:
-            QasmModule: The module with the measurements removed if in_place is False
+            QasmModule: The module with the measurements removed if in_place is False.
         """
         # copy first: the filtering rewrites nested box and if bodies in place, so it
         # has to run on the module that is being returned
@@ -277,11 +278,8 @@ def remove_measurements(self, in_place: bool = True) -> Optional["QasmModule"]:
     def has_barriers(self) -> bool:
         """Check if the module has any barrier operations.
 
-        Args:
-            None
-
         Returns:
-            bool: True if the module has barrier operations, False otherwise
+            bool: True if the module has barrier operations, False otherwise.
         """
         if self._has_barriers is None:
             self._has_barriers = False
@@ -300,7 +298,7 @@ def has_barriers(self) -> bool:
 
     @track_user_operation
     def remove_barriers(self, in_place: bool = True) -> Optional["QasmModule"]:
-        """Remove the barrier operations
+        """Remove the barrier operations.
 
         Args:
             in_place (bool): Flag to indicate if the removal should be done in place.
@@ -327,14 +325,14 @@ def remove_barriers(self, in_place: bool = True) -> Optional["QasmModule"]:
         return curr_module
 
     @track_user_operation
-    def remove_includes(self, in_place=True) -> Optional["QasmModule"]:
-        """Remove the include statements from the module
+    def remove_includes(self, in_place: bool = True) -> Optional["QasmModule"]:
+        """Remove the include statements from the module.
 
         Args:
             in_place (bool): Flag to indicate if the removal should be done in place.
 
         Returns:
-            QasmModule: The module with the includes removed if in_place is False, None otherwise
+            QasmModule: The module with the includes removed if in_place is False, None otherwise.
         """
         stmt_list = (
             self._statements
@@ -354,16 +352,16 @@ def remove_includes(self, in_place=True) -> Optional["QasmModule"]:
         return curr_module
 
     @track_user_operation
-    def depth(self, decompose_native_gates=True):
+    def depth(self, decompose_native_gates: bool = True) -> int:
         """Calculate the depth of the unrolled openqasm program.
 
         Args:
-        decompose_native_gates (bool): If True, calculate depth after decomposing gates.
-                                If False, treat all decompsable gates as a single gate operation.
-                                Defaults to True.
+            decompose_native_gates (bool): If True, calculate depth after decomposing gates.
+                If False, treat all decomposable gates as a single gate operation.
+                Defaults to True.
 
         Returns:
-            int: The depth of the current "unrolled" openqasm program
+            int: The depth of the current "unrolled" openqasm program.
         """
         # 1. Since the program will be unrolled before its execution on a QC, it makes sense to
         # calculate the depth of the unrolled program.
@@ -392,16 +390,22 @@ def depth(self, decompose_native_gates=True):
         max_depth = max(max_qubit_depth, max_clbit_depth)
         return max_depth
 
-    def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]):
+    def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]) -> None:
         """Remap the qubits in a register after removing idle qubits and update the operations
-        using this register accordingly"""
+        using this register accordingly.
+
+        Args:
+            reg_name (str): The name of the register to be remapped.
+            size (int): The size of the register.
+            idle_indices (list[int]): Indices of idle qubits to be remapped away.
+        """
 
         used_indices = [idx for idx in range(size) if idx not in idle_indices]
         new_size = size - len(idle_indices)
         idx_map = {used_indices[i]: i for i in range(new_size)}  # old_idx : new_idx
 
         # Example -
-        # reg_name = "q", original_size = 5, idle_indices = [1, 3]
+        # reg_name = "q", size = 5, idle_indices = [1, 3]
         # used_indices = [0, 2, 4], new_size = 3
         # idx_map = {0: 0, 2: 1, 4: 2}
 
@@ -440,11 +444,11 @@ def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]):
                     index_node.value = idx_map[index_node.value]  # type: ignore[union-attr]
 
     def _get_idle_qubit_indices(self) -> dict[str, list[int]]:
-        """Get the indices of the idle qubits in the module
+        """Get the indices of the idle qubits in the module.
 
         Returns:
-            dict[str, list[int]]: A dictionary mapping the register name to the list of idle qubit
-                                  indices in that register
+            dict[str, list[int]]: A dictionary mapping the register name
+                to the list of idle qubit indices in that register.
         """
         idle_qubits = [qubit for qubit in self._qubit_depths.values() if qubit.is_idle()]
 
@@ -457,10 +461,10 @@ def _get_idle_qubit_indices(self) -> dict[str, list[int]]:
 
         return qubit_indices
 
-    def populate_idle_qubits(self, in_place: bool = True):
-        """Populate the idle qubits in the module with identity gates
+    def populate_idle_qubits(self, in_place: bool = True) -> QasmModule:
+        """Populate the idle qubits in the module with identity gates.
 
-        Note: unrolling is not performed while calling this function
+        Note: unrolling is not performed while calling this function.
 
         Args:
             in_place (bool): Flag to indicate if the population should be done in place.
@@ -500,7 +504,7 @@ def populate_idle_qubits(self, in_place: bool = True):
 
         return qasm_module
 
-    def remove_idle_qubits(self, in_place: bool = True):
+    def remove_idle_qubits(self, in_place: bool = True) -> QasmModule:
         """Remove idle qubits from the module. Either collapse the size of a partially used
         quantum register OR remove the unused quantum register entirely.
 
@@ -553,7 +557,7 @@ def remove_idle_qubits(self, in_place: bool = True):
         return qasm_module
 
     @track_user_operation
-    def reverse_qubit_order(self, in_place=True):
+    def reverse_qubit_order(self, in_place: bool = True) -> QasmModule:
         """Reverse the order of qubits in the module.
 
         Will unroll the module if not already done.
@@ -589,7 +593,8 @@ def reverse_qubit_order(self, in_place=True):
                 if not isinstance(bit, qasm3_ast.IndexedIdentifier):
                     continue  # physical qubit ("$n"): not part of any register
                 curr_reg_name = bit.name.name
-                curr_reg_idx = bit.indices[0][0].value
+                index_node = bit.indices[0][0]  # type: ignore[index]
+                curr_reg_idx = index_node.value  # type: ignore[union-attr]
                 new_reg_idx = new_qubit_mappings[curr_reg_name][curr_reg_idx]
 
                 # make the idx -ve so that this is not touched
@@ -597,7 +602,7 @@ def reverse_qubit_order(self, in_place=True):
 
                 # idx -> -1 * idx - 1 as we also have to look at index 0
                 # which will remain 0 if we just multiply by -1
-                bit.indices[0][0].value = -1 * new_reg_idx - 1
+                index_node.value = -1 * new_reg_idx - 1  # type: ignore[union-attr]
 
         # remove the -ve marker
         for operation in iter_quantum_statements(qasm_module._unrolled_ast.statements):
@@ -605,9 +610,10 @@ def reverse_qubit_order(self, in_place=True):
             for bit in bit_list:
                 if not isinstance(bit, qasm3_ast.IndexedIdentifier):
                     continue
-                if bit.indices[0][0].value < 0:
-                    bit.indices[0][0].value += 1
-                    bit.indices[0][0].value *= -1
+                index_node = bit.indices[0][0]  # type: ignore[index]
+                if index_node.value < 0:  # type: ignore[union-attr]
+                    index_node.value += 1  # type: ignore[union-attr]
+                    index_node.value *= -1  # type: ignore[union-attr]
 
         # 3. update the original AST with the unrolled AST
         qasm_module._statements = qasm_module._unrolled_ast.statements
@@ -616,8 +622,8 @@ def reverse_qubit_order(self, in_place=True):
         return qasm_module
 
     @track_user_operation
-    def validate(self):
-        """Validate the module"""
+    def validate(self) -> None:
+        """Validate the module."""
         if self._validated_program is True:
             return
         try:
@@ -637,7 +643,7 @@ def validate(self):
         self._validated_program = True
 
     @track_user_operation
-    def unroll(self, **kwargs):
+    def unroll(self, **kwargs: Any) -> None:
         """Unroll the module into basic qasm operations.
 
         Args:
@@ -646,9 +652,9 @@ def unroll(self, **kwargs):
                 unroll_barriers (bool): If True, barriers will be unrolled. Defaults to True.
                 max_loop_iters (int): Max number of iterations for unrolling loops. Defaults to 1e9.
                 check_only (bool): If True, only check the program without executing it.
-                                   Defaults to False.
+                    Defaults to False.
                 consolidate_qubits (bool): If True, consolidate all quantum registers into
-                                           single register.
+                    single register.
 
         Raises:
             ValidationError: If the module fails validation during unrolling.
@@ -681,13 +687,13 @@ def unroll(self, **kwargs):
             raise err
 
     @track_user_operation
-    def rebase(self, target_basis_set, in_place=True):
+    def rebase(self, target_basis_set: BasisSet, in_place: bool = True) -> QasmModule:
         """Rebase the AST to use a specified target basis set.
 
-        Will unroll the module if not already done.
+        Note: Will unroll the module if not already done.
 
         Args:
-            target_basis_set: The target basis set to rebase the module to.
+            target_basis_set (BasisSet): The target basis set to rebase the module to.
             in_place (bool): Flag to indicate if the rebase operation should be done in place.
 
         Returns:
@@ -744,7 +750,7 @@ def _get_gate_counts(self) -> dict[str, int]:
         ]
         return dict(Counter(gate.name.name for gate in gate_nodes))
 
-    def compare(self, other_module: QasmModule):
+    def compare(self, other_module: QasmModule) -> None:
         """Compare two QasmModule objects across multiple attributes.
 
         Args:
@@ -768,7 +774,7 @@ def compare(self, other_module: QasmModule):
         self_ext_gates_str = "\n".join(self._external_gates)
         other_ext_gates_str = "\n".join(other_module._external_gates)
 
-        table_data = [
+        table_data: list[list[Any]] = [
             ["Qubits", self.num_qubits, other_module.num_qubits],
             ["Classical Bits", self.num_clbits, other_module.num_clbits],
             ["Measurements", self.has_measurements(), other_module.has_measurements()],
@@ -793,28 +799,28 @@ def compare(self, other_module: QasmModule):
         )
 
     def __str__(self) -> str:
-        """Return the string representation of the QASM program
+        """Return the string representation of the QASM program.
 
         Returns:
-            str: The string representation of the module
+            str: The string representation of the module.
         """
 
         if len(self._unrolled_ast.statements) > 1:
             return self._qasm_ast_to_str(self.unrolled_ast)
         return self._qasm_ast_to_str(self.original_program)
 
-    def copy(self):
-        """Return a deep copy of the module"""
+    def copy(self) -> QasmModule:
+        """Return a deep copy of the module."""
         return deepcopy(self)
 
     @abstractmethod
-    def _qasm_ast_to_str(self, qasm_ast):
-        """Convert the qasm AST to a string"""
+    def _qasm_ast_to_str(self, qasm_ast: Program) -> str:
+        """Convert the qasm AST to a string."""
 
     @abstractmethod
-    def accept(self, visitor):
-        """Accept a visitor for the mßodule
+    def accept(self, visitor: QasmVisitor) -> None:
+        """Accept a visitor for the module.
 
         Args:
-            visitor (QasmVisitor): The visitor to accept
+            visitor (QasmVisitor): The visitor to accept.
         """
diff --git a/src/pyqasm/modules/qasm2.py b/src/pyqasm/modules/qasm2.py
index bba821e..8a268ba 100644
--- a/src/pyqasm/modules/qasm2.py
+++ b/src/pyqasm/modules/qasm2.py
@@ -24,7 +24,7 @@
 from openqasm3.printer import dumps
 
 from pyqasm.exceptions import ValidationError, raise_qasm3_error
-from pyqasm.modules.base import QasmModule
+from pyqasm.modules.base import QasmModule, QasmVisitor
 from pyqasm.modules.qasm3 import Qasm3Module
 
 # the QASM 2.0 <qop> production: a gate application, a measurement or a reset.
@@ -54,7 +54,7 @@ class Qasm2Module(QasmModule):
         statements (list[Statement]): list of openqasm2 Statements.
     """
 
-    def __init__(self, name: str, program: Program):
+    def __init__(self, name: str, program: Program) -> None:
         super().__init__(name, program)
         self._unrolled_ast = Program(statements=[], version="2.0")
         self._whitelist_statements = {
@@ -70,8 +70,8 @@ def __init__(self, name: str, program: Program):
             qasm3_ast.QuantumBarrier,
         }
 
-    def _filter_statements(self):
-        """Filter statements according to the whitelist"""
+    def _filter_statements(self) -> None:
+        """Filter statements according to the whitelist."""
         for stmt in self._statements:
             stmt_type = type(stmt)
             if stmt_type not in self._whitelist_statements:
@@ -80,7 +80,7 @@ def _filter_statements(self):
                 self._filter_branch_body(stmt)
             # TODO: add more filtering here if needed
 
-    def _filter_branch_body(self, statement: qasm3_ast.BranchingStatement):
+    def _filter_branch_body(self, statement: qasm3_ast.BranchingStatement) -> None:
         """Filter the body of a conditional against what QASM 2.0 allows there.
 
         The QASM 2.0 grammar admits only a ``<qop>`` as the body of an ``if`` --
@@ -114,7 +114,7 @@ def _filter_branch_body(self, statement: qasm3_ast.BranchingStatement):
                 span=inner_stmt.span,
             )
 
-    def _format_declarations(self, qasm_str):
+    def _format_declarations(self, qasm_str: str) -> str:
         """Format the unrolled qasm for declarations in openqasm 2.0 format"""
         for declaration_type, replacement_type in [("qubit", "qreg"), ("bit", "creg")]:
             pattern = rf"{declaration_type}\[(\d+)\]\s+(\w+);"
@@ -122,20 +122,19 @@ def _format_declarations(self, qasm_str):
             qasm_str = re.sub(pattern, replacement, qasm_str)
         return qasm_str
 
-    def _qasm_ast_to_str(self, qasm_ast):
-        """Convert the qasm AST to a string"""
+    def _qasm_ast_to_str(self, qasm_ast: Program) -> str:
+        """Convert the qasm AST to a string."""
         # set the version to 2.0
         qasm_ast.version = "2.0"
         raw_qasm = dumps(qasm_ast, old_measurement=True)
         return self._format_declarations(raw_qasm)
 
     def to_qasm3(self, as_str: bool = False) -> str | Qasm3Module:
-        """Convert the module to openqasm3 format
+        """Convert the module to openqasm3 format.
 
         Args:
             as_str (bool): Flag to indicate if the conversion should be to a string
-                           or to a Qasm3Module object.
-                           Default is False.
+                or to a Qasm3Module object. Default is False.
 
         Returns:
             str | Qasm3Module: The module in openqasm3 format.
@@ -149,14 +148,14 @@ def to_qasm3(self, as_str: bool = False) -> str | Qasm3Module:
         qasm_program.version = "3.0"
         return dumps(qasm_program) if as_str else Qasm3Module(self._name, qasm_program)
 
-    def accept(self, visitor):
-        """Accept a visitor for the module
+    def accept(self, visitor: QasmVisitor) -> None:
+        """Accept a visitor for the module.
 
         Args:
-            visitor (QasmVisitor): The visitor to accept
+            visitor (QasmVisitor): The visitor to accept.
         """
         self._filter_statements()
         unrolled_stmt_list = visitor.visit_basic_block(self._statements)
         final_stmt_list = visitor.finalize(unrolled_stmt_list)
 
-        self.unrolled_ast.statements = final_stmt_list
+        self.unrolled_ast.statements = final_stmt_list  # type: ignore[assignment]
diff --git a/src/pyqasm/modules/qasm3.py b/src/pyqasm/modules/qasm3.py
index d440fb2..41e3458 100644
--- a/src/pyqasm/modules/qasm3.py
+++ b/src/pyqasm/modules/qasm3.py
@@ -22,7 +22,7 @@
 from openqasm3.ast import Pragma, Program, QASMNode
 from openqasm3.printer import Printer, PrinterState
 
-from pyqasm.modules.base import QasmModule
+from pyqasm.modules.base import QasmModule, QasmVisitor
 
 
 class Qasm3Printer(Printer):
@@ -70,23 +70,23 @@ class Qasm3Module(QasmModule):
         statements (list[Statement]): list of openqasm3 Statements.
     """
 
-    def __init__(self, name: str, program: Program):
+    def __init__(self, name: str, program: Program) -> None:
         super().__init__(name, program)
         self._unrolled_ast = Program(statements=[], version="3.0")
 
-    def _qasm_ast_to_str(self, qasm_ast):
-        """Convert the qasm AST to a string"""
+    def _qasm_ast_to_str(self, qasm_ast: Program) -> str:
+        """Convert the qasm AST to a string."""
         # set the version to 3.0
         qasm_ast.version = "3.0"
         return dumps(qasm_ast)
 
-    def accept(self, visitor):
-        """Accept a visitor for the module
+    def accept(self, visitor: QasmVisitor) -> None:
+        """Accept a visitor for the module.
 
         Args:
-            visitor (QasmVisitor): The visitor to accept
+            visitor (QasmVisitor): The visitor to accept.
         """
         unrolled_stmt_list = visitor.visit_basic_block(self._statements)
         final_stmt_list = visitor.finalize(unrolled_stmt_list)
 
-        self._unrolled_ast.statements = final_stmt_list
+        self._unrolled_ast.statements = final_stmt_list  # type: ignore[assignment]
diff --git a/src/pyqasm/pulse/visitor.py b/src/pyqasm/pulse/visitor.py
index 9fa5cc1..996a076 100644
--- a/src/pyqasm/pulse/visitor.py
+++ b/src/pyqasm/pulse/visitor.py
@@ -21,7 +21,7 @@
 """
 
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 
 import openqasm3.ast as qasm3_ast
 
@@ -754,14 +754,16 @@ def _visit_function_call(  # pylint: disable=too-many-branches, too-many-stateme
 
         return _return_value, [statement]
 
-    def visit_statement(self, statement: qasm3_ast.Statement) -> list[qasm3_ast.Statement]:
+    def visit_statement(
+        self, statement: qasm3_ast.Statement | qasm3_ast.Pragma
+    ) -> list[qasm3_ast.Statement]:
         """Visit a statement element.
 
         Args:
-            statement (qasm3_ast.Statement): The statement to visit.
+            statement (Statement | Pragma): The statement to visit.
 
         Returns:
-            None
+            list[Statement]: The list of resulting statements.
         """
         logger.debug("Visiting statement '%s'", str(statement))
         result = []
@@ -809,13 +811,13 @@ def visit_statement(self, statement: qasm3_ast.Statement) -> list[qasm3_ast.Stat
 
     def visit_basic_block(
         self,
-        stmt_list: list[qasm3_ast.Statement],
+        stmt_list: Sequence[qasm3_ast.Statement | qasm3_ast.Pragma],
         is_def_cal: bool,
     ) -> list[qasm3_ast.Statement]:
         """Visit a basic block of statements.
 
         Args:
-            stmt_list (list[qasm3_ast.Statement]): The list of statements to visit.
+            stmt_list (Sequence[Statement | Pragma]): The list of statements to visit.
             is_def_cal (bool): is the given statements from def_cal block.
 
         Returns:
diff --git a/src/pyqasm/visitor.py b/src/pyqasm/visitor.py
index ea9f5f3..62f52fa 100644
--- a/src/pyqasm/visitor.py
+++ b/src/pyqasm/visitor.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict, deque
 from functools import partial
 from io import StringIO
-from typing import Any, Callable, Optional, Sequence, Union, cast
+from typing import Any, Callable, Optional, Sequence, cast
 
 import numpy as np
 import openqasm3.ast as qasm3_ast
@@ -208,7 +208,8 @@ def _visit_quantum_register(
             register (QubitDeclaration): The register name and size.
 
         Returns:
-            None
+            list[QubitDeclaration]: The list containing the register,
+                or an empty list if self._check_only is True.
         """
         logger.debug("Visiting register '%s'", str(register))
 
@@ -289,15 +290,16 @@ def _get_op_bits(
         operation: Any,
         qubits: bool = True,
         function_qubit_sizes: Optional[dict[str, int]] = None,
-    ) -> list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]]:
+    ) -> list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier]:
         """Get the quantum / classical bits for the operation.
         Args:
             operation (Any): The operation to get qubits for.
             qubits (bool): Whether the bits are quantum bits or classical bits. Defaults to True.
         Returns:
-            The quantum or classical bits for the operation.
+            list[IndexedIdentifier | Identifier]: The quantum or classical bits for the operation,
+                or an empty list if check_only is true.
         """
-        openqasm_bits: list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]] = []
+        openqasm_bits: list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier] = []
         bit_list = []
 
         if isinstance(operation, qasm3_ast.QuantumMeasurementStatement):
@@ -414,7 +416,7 @@ def _check_variable_type_size(
             base_type (Any): Base type of the variable.
             is_const (bool): whether the statement is constant declaration or not.
         Returns:
-            Int: size of the variable base type.
+            int: size of the variable base type.
         """
         base_size = 1
         if not isinstance(base_type, qasm3_ast.BoolType):
@@ -464,8 +466,6 @@ def _check_variable_cast_type(
             base_type (Any): Base type of the declaration variable.
             base_size(Any): literal to get the base size of the declaration variable.
             is_const (bool): whether the statement is constant declaration or not.
-        Returns:
-            None
         """
         if not val_type:
             val_type = base_type
@@ -492,12 +492,17 @@ def _qubit_register_consolidation(
         Consolidate all quantum registers into a single register '__PYQASM_QUBITS__'.
 
         Args:
-            unrolled_stmts (list): The list of statements to process and modify in-place.
+            unrolled_stmts (list): The list of non-QubitDeclaration statements to append
+                after consolidating the quantum registers.
+            total_qubits (int): The total number of allocated qubits in quantum registers.
+
+        Returns:
+            A new list of statements with the quantum register statements consolidated.
 
         Raises:
             ValidationError: If the total number of qubits exceeds the available device qubits,
-                             or if the reserved register '__PYQASM_QUBITS__' is already declared
-                             in the original QASM program.
+                or if the reserved register '__PYQASM_QUBITS__' is already declared
+                in the original QASM program.
         """
         if total_qubits > self._module._device_qubits:  # type: ignore
             raise_qasm3_error(
@@ -526,13 +531,17 @@ def _qubit_register_consolidation(
         return _valid_statements
 
     def _handle_function_init_expression(
-        self, expression: Any, init_value: Any
+        self, expression: qasm3_ast.FunctionCall, init_value: Any
     ) -> None | qasm3_ast.Expression:
         """Handle function initialization expression.
 
         Args:
-            statement (Any): The statement to handle function initialization expression.
+            expression (FunctionCall): The statement to handle function initialization expression.
             init_value (Any): The value to handle function initialization expression.
+
+        Returns:
+            None | Expression: The resultant expression if
+                the expression is applied, otherwise None.
         """
         if isinstance(expression, qasm3_ast.FunctionCall):
             func_name = expression.name.name
@@ -570,10 +579,11 @@ def _visit_measurement(  # pylint: disable=too-many-locals,too-many-branches,too
         """Visit a measurement statement element.
 
         Args:
-            statement (qasm3_ast.QuantumMeasurementStatement): The measurement statement to visit.
+            statement (QuantumMeasurementStatement): The measurement statement to visit.
 
         Returns:
-            None
+            list[QuantumMeasurementStatement]: The list of unrolled measurements,
+                or an empty list if self._check_only is True.
         """
         logger.debug("Visiting measurement statement '%s'", str(statement))
 
@@ -611,7 +621,7 @@ def _visit_measurement(  # pylint: disable=too-many-locals,too-many-branches,too
                     )
             if is_pulse_gate:
                 return [statement]
-        # # TODO: handle in-function measurements
+        # TODO: handle in-function measurements
         source_name: str = (
             source.name if isinstance(source, qasm3_ast.Identifier) else source.name.name
         )
@@ -716,7 +726,7 @@ def _resolve_unindexed_reset_qubit(self, statement: qasm3_ast.QuantumReset) -> b
         register slot: a physical qubit ("$n") or the internal pulse register.
 
         Args:
-            statement (qasm3_ast.QuantumReset): The reset statement whose operand
+            statement (QuantumReset): The reset statement whose operand
                 is being resolved. Renamed in place for OpenPulse programs.
 
         Returns:
@@ -745,10 +755,11 @@ def _visit_reset(self, statement: qasm3_ast.QuantumReset) -> list[qasm3_ast.Quan
         """Visit a reset statement element.
 
         Args:
-            statement (qasm3_ast.QuantumReset): The reset statement to visit.
+            statement (QuantumReset): The reset statement to visit.
 
         Returns:
-            None
+            list[QuantumReset]: The list of unrolled resets,
+                or an empty list if self._check_only is True.
         """
         logger.debug("Visiting reset statement '%s'", str(statement))
         if self._resolve_unindexed_reset_qubit(statement):
@@ -806,7 +817,15 @@ def _expand_barrier_ranges(
     ) -> list:
         """Replace RangeDefinition-containing qubits in a barrier with their
         expanded IndexedIdentifier equivalents so that consolidate_qubit_registers
-        only sees IntegerLiteral indices."""
+        only sees IntegerLiteral indices.
+
+        Args:
+            barrier (QuantumBarrier): The barrier whose qubits are being expanded.
+            barrier_qubits (list[IndexedIdentifier | Identifier]): The expanded qubit operands.
+
+        Returns:
+            list: The resultant list of consolidated qubits.
+        """
         consolidated_qubits: list = []
         expanded_idx = 0
         for op_qubit in barrier.qubits:
@@ -841,10 +860,11 @@ def _visit_barrier(  # pylint: disable=too-many-locals, too-many-branches
         """Visit a barrier statement element.
 
         Args:
-            statement (qasm3_ast.QuantumBarrier): The barrier statement to visit.
+            statement (QuantumBarrier): The barrier statement to visit.
 
         Returns:
-            None
+            list[QuantumBarrier]: The list containing a single multi-qubit barrier statement,
+                or an empty list if self._check_only is True.
         """
         valid_open_pulse_qubits = False
         for op_qubit in barrier.qubits:
@@ -945,7 +965,7 @@ def _get_op_parameters(self, operation: qasm3_ast.QuantumGate) -> list[float]:
         """Get the parameters for the operation.
 
         Args:
-            operation (qasm3_ast.QuantumGate): The operation to get parameters for.
+            operation (QuantumGate): The operation to get parameters for.
 
         Returns:
             list[float]: The parameters for the operation.
@@ -972,7 +992,7 @@ def _visit_gate_definition(self, definition: qasm3_ast.QuantumGateDefinition) ->
             definition (qasm3_ast.QuantumGateDefinition): The gate definition to visit.
 
         Returns:
-            None
+            An empty list.
         """
         gate_name = definition.name.name
         if gate_name in self._custom_gates:
@@ -987,16 +1007,17 @@ def _visit_gate_definition(self, definition: qasm3_ast.QuantumGateDefinition) ->
 
     def _unroll_multiple_target_qubits(
         self, operation: qasm3_ast.QuantumGate, gate_qubit_count: int
-    ) -> list[list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]]]:
+    ) -> list[list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier]]:
         """Unroll the complete list of all qubits that the given operation is applied to.
-           E.g. this maps 'cx q[0], q[1], q[2], q[3]' to [[q[0], q[1]], [q[2], q[3]]]
+            E.g. this maps 'cx q[0], q[1], q[2], q[3]' to [[q[0], q[1]], [q[2], q[3]]]
 
         Args:
-            operation (qasm3_ast.QuantumGate): The gate to be applied.
+            operation (QuantumGate): The gate to be applied.
             gate_qubit_count (list[int]): The number of qubits that a single gate acts on.
 
         Returns:
-            The list of all targets that the unrolled gate should act on.
+            list[list[IndexedIdentifier | Identifier]]: The list of all targets that
+                the unrolled gate should act on.
         """
         op_qubits = self._get_op_bits(operation, qubits=True)
         if len(op_qubits) <= 0 or len(op_qubits) % gate_qubit_count != 0:
@@ -1014,20 +1035,21 @@ def _unroll_multiple_target_qubits(
     def _broadcast_gate_operation(
         self,
         gate_function: Callable,
-        all_targets: list[list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]]],
+        all_targets: list[list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier]],
         ctrls: Optional[list[qasm3_ast.IndexedIdentifier]] = None,
     ) -> list[qasm3_ast.QuantumGate]:
         """Broadcasts the application of a gate onto multiple sets of target qubits.
 
         Args:
-            gate_function (callable): The gate that should be applied to multiple target qubits.
+            gate_function (Callable): The gate that should be applied to multiple target qubits.
             (All arguments of the callable should be qubits, i.e. all non-qubit arguments of the
             gate should already be evaluated, e.g. using functools.partial).
-            all_targets (list[list[qasm3_ast.IndexedIdentifier]]):
+            all_targets (list[list[IndexedIdentifier]]):
                 The list of target qubits.
                 The length of this list indicates the number of time the gate is invoked.
+
         Returns:
-            List of all executed gates.
+            list[QuantumGate]: The list of all executed gates.
         """
         result = []
         if ctrls is None:
@@ -1040,10 +1062,10 @@ def _register_physical_qubit(self, name: str) -> int:
         """Register a physical qubit ``$n`` for depth / count tracking if not already known.
 
         Args:
-            name: The physical qubit identifier string (e.g. ``"$0"``).
+            name (str): The physical qubit identifier string (e.g. ``"$0"``).
 
         Returns:
-            The physical qubit index.
+            int: The physical qubit index.
         """
         phys_idx = int(name[1:])
         if (name, phys_idx) not in self._module._qubit_depths:
@@ -1074,6 +1096,12 @@ def _get_qubit_name_and_id(
         Physical qubits are represented as ``Identifier("$n")`` and carry their
         index in the name itself.  Virtual qubits are ``IndexedIdentifier`` with
         an explicit index in ``.indices``.
+
+        Args:
+            qubit (IndexedIdentifier | Identifier): The qubit to get the name and id of.
+
+        Returns:
+            tuple[str, int]: A tuple containing the name and id of the qubit.
         """
         if isinstance(qubit, qasm3_ast.Identifier):
             # Physical qubit: name is "$n", index is n.
@@ -1084,16 +1112,15 @@ def _get_qubit_name_and_id(
 
     def _update_qubit_depth_for_gate(
         self,
-        all_targets: list[list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]]],
+        all_targets: list[list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier]],
         ctrls: list[qasm3_ast.IndexedIdentifier],
-    ):
+    ) -> None:
         """Updates the depth of the circuit after applying a broadcasted gate.
 
         Args:
-            all_targes: The list of qubits on which a gate was just added.
-
-        Returns:
-            None
+            all_targets (list[list[IndexedIdentifier | Identifier]]):
+                The list of qubits on which a gate was just added.
+            ctrls (list[IndexedIdentifier]): The list of control qubits for the gate.
         """
         if not self._recording_ext_gate_depth:
             for qubit_subset in all_targets:
@@ -1119,20 +1146,17 @@ def _visit_basic_gate_operation(
         """Visit a gate operation element.
 
         Args:
-            operation (qasm3_ast.QuantumGate): The gate operation to visit.
+            operation (QuantumGate): The gate operation to visit.
             inverse (bool): Whether the operation is an inverse operation. Defaults to False.
-
-                          - if inverse is True, we apply check for different cases in the
-                            map_qasm_inv_op_to_callable method.
-
-                          - Only rotation and S / T gates are affected by this inversion. For S/T
-                            gates we map them to Sdg / Tdg and vice versa.
-
-                          - For rotation gates, we map to the same gates but invert the rotation
-                            angles.
+                    If inverse is True, we apply check for different cases in the
+                    map_qasm_inv_op_to_callable method.
+                    Only rotation and S / T gates are affected by this inversion. For S/T
+                    gates we map them to Sdg / Tdg and vice versa.
+                    For rotation gates, we map to the same gates but invert the rotation angles.
 
         Returns:
-            None
+            list[QuantumGate]: The list of gates after unrolling,
+                or an empty list if self._check_only is True.
 
         Raises:
             ValidationError: If the number of qubits is invalid.
@@ -1237,24 +1261,24 @@ def _visit_custom_gate_operation(
         """Visit a custom gate operation element recursively.
 
         Args:
-            operation (qasm3_ast.QuantumGate): The gate operation to visit.
+            operation (QuantumGate): The gate operation to visit.
             inverse (bool): Whether the operation is an inverse operation. Defaults to False.
-
-                            If True, the gate operation is applied in reverse order and the
-                            inverse modifier is appended to each gate call.
-                            See https://openqasm.com/language/gates.html#inverse-modifier
-                            for more clarity.
+                If True, the gate operation is applied in reverse order and the
+                inverse modifier is appended to each gate call.
+                See https://openqasm.com/language/gates.html#inverse-modifier
+                for more clarity.
 
         Returns:
-            None
+            list[QuantumGate | QuantumPhase]: The list of gates and phase operations
+                after unrolling, or an empty list if self._check_only is True.
         """
         logger.debug("Visiting custom gate operation '%s'", str(operation))
         if ctrls is None:
             ctrls = []
         gate_name: str = operation.name.name
         gate_definition: qasm3_ast.QuantumGateDefinition = self._custom_gates[gate_name]
-        op_qubits: list[Union[qasm3_ast.IndexedIdentifier, qasm3_ast.Identifier]] = (
-            self._get_op_bits(operation, qubits=True)
+        op_qubits: list[qasm3_ast.IndexedIdentifier | qasm3_ast.Identifier] = self._get_op_bits(
+            operation, qubits=True
         )
 
         Qasm3Validator.validate_gate_call(operation, gate_definition, len(op_qubits))
@@ -1344,14 +1368,14 @@ def _visit_external_gate_operation(
         Args:
             operation (qasm3_ast.QuantumGate): The external gate operation to visit.
             inverse (bool): Whether the operation is an inverse operation. Defaults to False.
-
-                            If True, the gate operation is applied in reverse order and the
-                            inverse modifier is appended to each gate call.
-                            See https://openqasm.com/language/gates.html#inverse-modifier
-                            for more clarity.
+                If True, the gate operation is applied in reverse order and the
+                inverse modifier is appended to each gate call.
+                See https://openqasm.com/language/gates.html#inverse-modifier
+                for more clarity.
 
         Returns:
-            list[qasm3_ast.QuantumGate]: The quantum gate that was collected.
+            list[QuantumGate]: The list containing the quantum gate that was collected,
+                or an empty list if self._check_only is True.
         """
         logger.debug("Visiting external gate operation '%s'", str(operation))
         gate_name: str = operation.name.name
@@ -1418,11 +1442,16 @@ def _visit_phase_operation(
         """Visit a phase operation element.
 
         Args:
-            operation (qasm3_ast.QuantumPhase): The phase operation to visit.
+            operation (QuantumPhase): The phase operation to visit.
             inverse (bool): Whether the operation is an inverse operation. Defaults to False.
+                If True, the gate operation is applied in reverse order and the
+                inverse modifier is appended to each gate call.
+                See https://openqasm.com/language/gates.html#inverse-modifier
+                for more clarity.
 
         Returns:
-            list[qasm3_ast.Statement]: The unrolled quantum phase operation.
+            list[QuantumPhase]: The unrolled quantum phase operation, or an empty list
+                if self._check_only is True.
         """
         logger.debug("Visiting phase operation '%s'", str(operation))
         if ctrls is None:
@@ -1476,10 +1505,13 @@ def _visit_generic_gate_operation(  # pylint: disable=too-many-branches, too-man
         """Visit a gate operation element.
 
         Args:
-            operation (qasm3_ast.QuantumGate): The gate operation to visit.
+            operation (QuantumGate | QuantumPhase): The gate operation to visit.
+            ctrls (list[IndexedIdentifier]): An optional list of control qubits
+                on the gate operation.
 
         Returns:
-            None
+            list[QuantumGate | QuantumPhase]: The list of gates and phase operations
+                after unrolling, or an empty list if self._check_only is True.
         """
         operation, ctrls = copy.copy(operation), copy.copy(ctrls)
         negctrls = []
@@ -1638,10 +1670,11 @@ def _visit_constant_declaration(
         type variables and not arrays. Assignment is mandatory in constant declaration.
 
         Args:
-            statement (qasm3_ast.ConstantDeclaration): The constant declaration to visit.
+            statement (ConstantDeclaration): The constant declaration to visit.
 
         Returns:
-            None
+            list[Statement]: The list containing the unrolled statement, or an empty
+                list if self._check_only is True.
         """
         statements = []
         var_name = statement.identifier.name
@@ -1757,10 +1790,11 @@ def _visit_classical_declaration(
         """Visit a classical operation element.
 
         Args:
-            statement (ClassicalType): The classical operation to visit.
+            statement (ClassicalDeclaration): The classical operation to visit.
 
         Returns:
-            None
+            list[Statement]: The list containing the unrolled statement, or an empty list
+                if self._check_only is True.
         """
         statements = []
         var_name = statement.identifier.name
@@ -1995,10 +2029,11 @@ def _visit_classical_assignment(
         """Visit a classical assignment element.
 
         Args:
-            statement (qasm3_ast.ClassicalAssignment): The classical assignment to visit.
+            statement (ClassicalAssignment): The classical assignment to visit.
 
         Returns:
-            list[qasm3_ast.Statement]: The list of statements generated by the assignment.
+            list[Statement]: The list containing the unrolled statement, or an empty list
+                if self._check_only is True.
         """
         statements = []
         lvalue = statement.lvalue
@@ -2170,7 +2205,7 @@ def _evaluate_array_initialization(
         """Evaluate an array initialization.
 
         Args:
-            array_literal (qasm3_ast.ArrayLiteral): The array literal to evaluate.
+            array_literal (ArrayLiteral): The array literal to evaluate.
             dimensions (list[int]): The dimensions of the array.
             base_type (Any): The base type of the array.
 
@@ -2212,10 +2247,11 @@ def _visit_branching_statement(
         """Visit a branching statement element.
 
         Args:
-            statement (qasm3_ast.BranchingStatement): The branching statement to visit.
+            statement (BranchingStatement): The branching statement to visit.
 
         Returns:
-            None
+            list[Statement]: The list of unrolled branch statements, or an empty list
+                if self._check_only is True.
         """
         self._scope_manager.push_context(Context.BLOCK)
         self._scope_manager.push_scope({})
@@ -2340,7 +2376,15 @@ def ravel(bit_ind):
         return result  # type: ignore[return-value]
 
     def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.Statement]:
-        # Compute loop variable values
+        """Visit a for-in loop statement element.
+
+        Args:
+            statement (ForInLoop): The for-in loop statement to visit.
+
+        Returns:
+            list[Statement]: The list containing the loop statements,
+                or an empty list if self._check_only is True.
+        """
         irange = []
         if isinstance(statement.set_declaration, qasm3_ast.RangeDefinition):
             init_exp = statement.set_declaration.start
@@ -2422,10 +2466,13 @@ def _visit_subroutine_definition(
            Reference: https://openqasm.com/language/subroutines.html#subroutines
 
         Args:
-            statement (qasm3_ast.SubroutineDefinition): The subroutine definition to visit.
+            statement (SubroutineDefinition | ExternDeclaration):
+                The subroutine definition to visit.
 
         Returns:
-            None
+            Sequence[None | ExternDeclaration]: The list containing the statement
+                if it is an ExternDeclaration, otherwise an empty list.
+                Returns an empty list if self._check_only is True.
         """
         fn_name = statement.name.name
         statements = []
@@ -2473,10 +2520,15 @@ def _visit_function_call(
         """Visit a function call element.
 
         Args:
-            statement (qasm3_ast.FunctionCall): The function call to visit.
-        Returns:
-            None
+            statement (FunctionCall): The function call to visit.
 
+        Returns:
+            tuple[Any | None, list[Statement | FunctionCall]]: If the function is undefined,
+                returns a tuple of None and an empty list.
+                If the function is defined, returns a tuple of the function's return
+                value and a list containing the function statement.
+                If self._check_only is True, instead return a tuple of
+                the function's return value and an empty list.
         """
         fn_name = statement.name.name
         if fn_name not in self._subroutine_defns and fn_name not in FUNCTION_MAP:
@@ -2595,12 +2647,16 @@ def _visit_while_loop(self, statement: qasm3_ast.WhileLoop) -> list[qasm3_ast.St
         """Visit a while-loop element.
 
         Args:
-            statement (qasm3_ast.WhileLoop) - the while-loop AST node
+            statement (WhileLoop): The while-loop AST node.
+
         Returns:
-            list[qasm3_ast.Statement] - flattened/unrolled statements
+            list[Statement]: The list of unrolled statements from the while-loop, or an
+                empty list if self._check_only is True.
+
         Raises:
-            ValidationError - if loop condition is non-classical or dynamic
-            LoopLimitExceededError - if the loop exceeds the maximum limit"""
+            ValidationError: If loop condition is non-classical or dynamic.
+            LoopLimitExceededError: If the loop exceeds the maximum limit.
+        """
 
         result = []
 
@@ -2652,10 +2708,10 @@ def _visit_alias_statement(self, statement: qasm3_ast.AliasStatement) -> list[No
         """Visit an alias statement element.
 
         Args:
-            statement (qasm3_ast.AliasStatement): The alias statement to visit.
+            statement (AliasStatement): The alias statement to visit.
 
         Returns:
-            None
+            list[None]: An empty list.
         """
         # pylint: disable=too-many-branches
         target = statement.target
@@ -2768,10 +2824,11 @@ def _visit_switch_statement(  # type: ignore[return]
         """Visit a switch statement element.
 
         Args:
-            statement (qasm3_ast.SwitchStatement): The switch statement to visit.
+            statement (SwitchStatement): The switch statement to visit.
 
         Returns:
-            list[qasm3_ast.Statement]: The list of statements generated by the switch statement.
+            list[Statement]: The list of statements generated by the switch statement,
+                or an empty list if self._check_only is True.
         """
         # 1. analyze the target - it should ONLY be int, not casted
         switch_target = statement.target
@@ -2852,12 +2909,18 @@ def _evaluate_case(statements):
             default_stmts = statement.default.statements
             return _evaluate_case(default_stmts)
 
-    def _resolve_duration_unit(self, time_var) -> qasm3_ast.TimeUnit:
+    def _resolve_duration_unit(self, time_var: qasm3_ast.Expression) -> qasm3_ast.TimeUnit:
         """Determine the output unit for a duration literal.
 
         `dt` is backend-dependent and not convertible to SI units without a known
         sample rate. Preserve it when the source unit was `dt` (or a device cycle
         time is set). SI units are already converted to ns by the evaluator.
+
+        Args:
+            time_var (Expression): A DurationLiteral to check.
+
+        Returns:
+            TimeUnit: The unit of the DurationLiteral.
         """
         source_is_dt = (
             isinstance(time_var, qasm3_ast.DurationLiteral)
@@ -2874,10 +2937,13 @@ def _visit_delay_statement(
     ) -> list[qasm3_ast.Statement]:
         """
         Visit a DelayInstruction statement.
+
         Args:
-            statement (qasm3_ast.DelayInstruction): The DelayInstruction statement to visit.
+            statement (DelayInstruction): The DelayInstruction statement to visit.
+
         Returns:
-            list[qasm3_ast.Statement]: The list of statements generated by the DelayInstruction.
+            list[Statement]: The list of statements generated by the DelayInstruction,
+                or an empty list if self._check_only is True.
         """
         _delay_time_var = statement.duration
         global_scope = self._scope_manager.get_global_scope()
@@ -2997,10 +3063,13 @@ def _visit_pragma(self, statement: qasm3_ast.Pragma) -> list[qasm3_ast.Pragma]:
     def _visit_box_statement(self, statement: qasm3_ast.Box) -> list[qasm3_ast.Statement]:
         """
         Visit a Box statement.
+
         Args:
-            statement (qasm3_ast.Box): The Box statement node to visit.
+            statement (Box): The Box statement node to visit.
+
         Returns:
-            list[qasm3_ast.Statement]: The list of statements generated by the Box statement.
+            list[Statement]: The list of unrolled statements from the Box statement, or an
+                empty list if self._check_only is True.
         """
         statements = []
         outer_verbatim = self._in_verbatim_box
@@ -3075,14 +3144,15 @@ def _visit_box_statement(self, statement: qasm3_ast.Box) -> list[qasm3_ast.State
 
     def _visit_calibration_definition(
         self, statement: qasm3_ast.CalibrationDefinition
-    ) -> list[Any]:
+    ) -> list[qasm3_ast.Statement]:
         """Visit a calibration definition element.
 
         Args:
-            statement (qasm3_ast.CalibrationDefinition): The calibration definition to visit.
+            statement (CalibrationDefinition): The calibration definition to visit.
 
         Returns:
-            None
+            list[Statement]: The list of unrolled statements, or the original statement
+                in a list if self._check_only is True.
         """
         from openpulse.parser import (  # pylint: disable=import-outside-toplevel
             OpenPulseParsingError,
@@ -3222,14 +3292,17 @@ def _visit_calibration_definition(
 
         return [statement]
 
-    def _visit_calibration_statement(self, statement: qasm3_ast.CalibrationStatement) -> list[Any]:
+    def _visit_calibration_statement(
+        self, statement: qasm3_ast.CalibrationStatement
+    ) -> list[qasm3_ast.Statement]:
         """Visit a calibration statement element.
 
         Args:
-            statement (qasm3_ast.CalibrationStatement): The calibration statement to visit.
+            statement (CalibrationStatement): The calibration statement to visit.
 
         Returns:
-            None
+            list[Statement]: The list of unrolled statements, or the original statement
+                in a list if self._check_only is True.
         """
         from openpulse.parser import (  # pylint: disable=import-outside-toplevel
             OpenPulseParsingError,
@@ -3295,10 +3368,10 @@ def _visit_calibration_grammar_declaration(
         """Visit a calibration grammar declaration element.
 
         Args:
-            statement (qasm3_ast.CalibrationGrammarDeclaration): The calibration grammar declaration
+            statement (CalibrationGrammarDeclaration): The calibration grammar declaration
 
         Returns:
-            None
+            list[Statement]: The list of unrolled statements.
         """
         if statement.name != "openpulse":
             raise_qasm3_error(
@@ -3315,10 +3388,11 @@ def _visit_include(self, include: qasm3_ast.Include) -> list[qasm3_ast.Statement
         """Visit an include statement element.
 
         Args:
-            include (qasm3_ast.Include): The include statement to visit.
+            include (Include): The include statement to visit.
 
         Returns:
-            None
+            list[Statement]: A list containing the include statement,
+                or an empty list if self._check_only is True.
         """
         filename = include.filename
         if filename in self._included_files:
@@ -3337,10 +3411,10 @@ def visit_statement(
         """Visit a statement element.
 
         Args:
-            statement (qasm3_ast.Statement | qasm3_ast.Pragma): The statement to visit.
+            statement (Statement | Pragma): The statement to visit.
 
         Returns:
-            None
+            list[Statement]: The list of unrolled statements.
         """
         logger.debug("Visiting statement '%s'", str(statement))
         result = []
@@ -3373,27 +3447,27 @@ def visit_basic_block(
         """Visit a basic block of statements.
 
         Args:
-            stmt_list (Sequence[qasm3_ast.Statement | qasm3_ast.Pragma]): The statements to visit.
+            stmt_list (Sequence[Statement | Pragma]): The list of statements to visit.
 
         Returns:
-            list[qasm3_ast.Statement]: The list of unrolled statements.
+            list[Statement]: The list of unrolled statements.
         """
         result = []
         for stmt in stmt_list:
             result.extend(self.visit_statement(stmt))
         return result
 
-    def finalize(self, unrolled_stmts):
+    def finalize(self, unrolled_stmts: list[qasm3_ast.Statement]) -> list[qasm3_ast.Statement]:
         """Finalize the unrolled statements.
         Rules:
         - Remove qubit args from phase operations if ALL qubits are used
         To add more rules if needed
 
         Args:
-            unrolled_stmts (list[qasm3_ast.Statement]): The list of unrolled statements.
+            unrolled_stmts (list[Statement]): The list of unrolled statements.
 
         Returns:
-            list[qasm3_ast.Statement]: The list of finalized statements.
+            list[Statement]: The list of finalized statements.
 
         """
         # remove the gphase qubits if they use ALL qubits