fix: Correct MitTask type

qermit/taskgraph/mittask.py

@@ -16,7 +16,7 @@
 from collections import namedtuple
 from enum import Enum
 from types import MethodType
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from pytket import Bit, Circuit, Qubit
 from pytket.backends import ResultHandle
@@ -55,6 +55,22 @@ class IOTask(Enum):
     Dict,
 ]
 
+WireList = Union[
+    List[CircuitShots],
+    List[Circuit],
+    List[BackendResult],
+    List[ResultHandle],
+    List[AnsatzCircuit],
+    List[ObservableExperiment],
+    List[int],
+    List[float],
+    List[bool],
+    List[str],
+    List[QubitPauliOperator],
+    List[Dict[Qubit, Bit]],
+    List[Dict],
+]
+
 
 class MitTask:
     """
@@ -98,7 +114,7 @@ def n_in_wires(self):
     def n_out_wires(self):
         return self._n_out_wires
 
-    def __call__(self, input_wires: List[Wire]) -> List[Wire]:
+    def __call__(self, input_wires: Tuple[WireList, ...]) -> Tuple[WireList, ...]:
         return self.run(*input_wires)
 
     def __str__(self):

tests/leakage_gadget_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import cast
+
 from pytket import Circuit
 from pytket.extensions.quantinuum.backends.leakage_gadget import (
     prune_shots_detected_as_leaky,
@@ -69,13 +71,15 @@ def test_compare_with_prune() -> None:
             ),
             n_shots=detection_circuit_shots.Shots,
         )
-        for detection_circuit_shots in detection_circuit_shots_list
+        for detection_circuit_shots in cast(
+            list[CircuitShots], detection_circuit_shots_list
+        )
     ]
 
     qermit_result_list = postselection_task(
         (
             result_list,
-            postselect_mgr_list,
+            postselect_mgr_list,  # type: ignore
         )
     )
     pytket_result_list = [