From e551c8d71857523f676d9bf7ccff38b04b57e3fd Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 10:37:08 +0000
Subject: [PATCH 1/9] fix: Correct MitTask type

---
 qermit/taskgraph/mittask.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 65d1c14a..67421b36 100644
--- a/qermit/taskgraph/mittask.py
+++ b/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
@@ -98,7 +98,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[List[Wire]]) -> Tuple[List[Wire]]:
         return self.run(*input_wires)
 
     def __str__(self):

From 06e14bed7a668bdbe4542a77a6fabfb11712a14e Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 11:57:02 +0000
Subject: [PATCH 2/9] Allow many values in tuple return

---
 qermit/taskgraph/mittask.py  | 4 +++-
 tests/leakage_gadget_test.py | 7 +++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 67421b36..16db2f60 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -22,6 +22,7 @@
 from pytket.backends import ResultHandle
 from pytket.backends.backendresult import BackendResult
 from pytket.utils import QubitPauliOperator
+from qermit.postselection.postselect_manager import PostselectMgr
 
 
 class IOTask(Enum):
@@ -40,6 +41,7 @@ class IOTask(Enum):
 )
 
 Wire = Union[
+    PostselectMgr,
     CircuitShots,
     Circuit,
     BackendResult,
@@ -98,7 +100,7 @@ def n_in_wires(self):
     def n_out_wires(self):
         return self._n_out_wires
 
-    def __call__(self, input_wires: Tuple[List[Wire]]) -> Tuple[List[Wire]]:
+    def __call__(self, input_wires: Tuple[List[Wire], ...]) -> Tuple[List[Wire], ...]:
         return self.run(*input_wires)
 
     def __str__(self):
diff --git a/tests/leakage_gadget_test.py b/tests/leakage_gadget_test.py
index a59f4f15..3081fac7 100644
--- a/tests/leakage_gadget_test.py
+++ b/tests/leakage_gadget_test.py
@@ -26,6 +26,9 @@
 from qermit.postselection.postselect_mitres import gen_postselect_task
 from qermit.taskgraph import gen_compiled_MitRes
 
+from typing import cast
+from qermit.postselection import PostselectMgr
+
 
 def test_leakage_gadget() -> None:
     backend = MockQuantinuumBackend()
@@ -69,13 +72,13 @@ 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,
+            cast(list[PostselectMgr], postselect_mgr_list),
         )
     )
     pytket_result_list = [

From 212d85048fc79b3f79542e9e3c0a36da1b03dddd Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 12:08:33 +0000
Subject: [PATCH 3/9] Allow lists of wire

---
 qermit/taskgraph/mittask.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 16db2f60..fe1bbe7b 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -57,6 +57,23 @@ class IOTask(Enum):
     Dict,
 ]
 
+WireList = Union[
+    List[PostselectMgr],
+    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:
     """
@@ -100,7 +117,7 @@ def n_in_wires(self):
     def n_out_wires(self):
         return self._n_out_wires
 
-    def __call__(self, input_wires: Tuple[List[Wire], ...]) -> Tuple[List[Wire], ...]:
+    def __call__(self, input_wires: Tuple[WireList, ...]) -> Tuple[WireList, ...]:
         return self.run(*input_wires)
 
     def __str__(self):

From e5b0f374dcba71acc76ed5ccbcadca4668a96da0 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 12:38:46 +0000
Subject: [PATCH 4/9] Remove PostselectMgr from typing

---
 qermit/taskgraph/mittask.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index fe1bbe7b..9e862527 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -22,7 +22,6 @@
 from pytket.backends import ResultHandle
 from pytket.backends.backendresult import BackendResult
 from pytket.utils import QubitPauliOperator
-from qermit.postselection.postselect_manager import PostselectMgr
 
 
 class IOTask(Enum):
@@ -41,7 +40,6 @@ class IOTask(Enum):
 )
 
 Wire = Union[
-    PostselectMgr,
     CircuitShots,
     Circuit,
     BackendResult,
@@ -58,7 +56,6 @@ class IOTask(Enum):
 ]
 
 WireList = Union[
-    List[PostselectMgr],
     List[CircuitShots],
     List[Circuit],
     List[BackendResult],

From e8247d748aa298d778ac7c289cbcb780184e3770 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 12:54:40 +0000
Subject: [PATCH 5/9] Use ignore in the case of postselect manager

---
 qermit/taskgraph/mittask.py  | 2 +-
 tests/leakage_gadget_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 9e862527..e7b22a86 100644
--- a/qermit/taskgraph/mittask.py
+++ b/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, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from pytket import Bit, Circuit, Qubit
 from pytket.backends import ResultHandle
diff --git a/tests/leakage_gadget_test.py b/tests/leakage_gadget_test.py
index 3081fac7..7a51449c 100644
--- a/tests/leakage_gadget_test.py
+++ b/tests/leakage_gadget_test.py
@@ -78,7 +78,7 @@ def test_compare_with_prune() -> None:
     qermit_result_list = postselection_task(
         (
             result_list,
-            cast(list[PostselectMgr], postselect_mgr_list),
+            postselect_mgr_list,  # type: ignore
         )
     )
     pytket_result_list = [

From ed084e7bfe6fd4973ee9b83cdc708387871b8a2d Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:07:13 +0000
Subject: [PATCH 6/9] Use type parameter

---
 .github/workflows/build_and_test.yml       | 4 ++++
 qermit/postselection/postselect_manager.py | 3 ++-
 qermit/taskgraph/mittask.py                | 5 +++++
 tests/leakage_gadget_test.py               | 3 +--
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index f02b7981..3167a138 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -86,6 +86,10 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v4
+    - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.12
     - name: Build qermit
       run: |
         pip install .[tests]
diff --git a/qermit/postselection/postselect_manager.py b/qermit/postselection/postselect_manager.py
index 2702a7b5..d2c9a460 100644
--- a/qermit/postselection/postselect_manager.py
+++ b/qermit/postselection/postselect_manager.py
@@ -4,9 +4,10 @@
 from pytket.backends.backendresult import BackendResult
 from pytket.circuit import Bit
 from pytket.utils.outcomearray import OutcomeArray
+from qermit.taskgraph.mittask import PostselectMgrType
 
 
-class PostselectMgr:
+class PostselectMgr[PostselectMgrType]:
     """Class for tracking and applying post selection to results.
 
     An example use case might be the following. Here a Bell state is
diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index e7b22a86..67d7531a 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -22,6 +22,7 @@
 from pytket.backends import ResultHandle
 from pytket.backends.backendresult import BackendResult
 from pytket.utils import QubitPauliOperator
+from typing import TypeVar
 
 
 class IOTask(Enum):
@@ -39,7 +40,10 @@ class IOTask(Enum):
     "ObservableExperiment", ["AnsatzCircuit", "ObservableTracker"]
 )
 
+PostselectMgrType = TypeVar('PostselectMgrType')
+
 Wire = Union[
+    PostselectMgrType,
     CircuitShots,
     Circuit,
     BackendResult,
@@ -56,6 +60,7 @@ class IOTask(Enum):
 ]
 
 WireList = Union[
+    List[PostselectMgrType],
     List[CircuitShots],
     List[Circuit],
     List[BackendResult],
diff --git a/tests/leakage_gadget_test.py b/tests/leakage_gadget_test.py
index 7a51449c..70096471 100644
--- a/tests/leakage_gadget_test.py
+++ b/tests/leakage_gadget_test.py
@@ -27,7 +27,6 @@
 from qermit.taskgraph import gen_compiled_MitRes
 
 from typing import cast
-from qermit.postselection import PostselectMgr
 
 
 def test_leakage_gadget() -> None:
@@ -78,7 +77,7 @@ def test_compare_with_prune() -> None:
     qermit_result_list = postselection_task(
         (
             result_list,
-            postselect_mgr_list,  # type: ignore
+            postselect_mgr_list,
         )
     )
     pytket_result_list = [

From 539e19297c9257af081eb39370f1c45d3908b5cc Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:09:18 +0000
Subject: [PATCH 7/9] Correct yml syntax

---
 .github/workflows/build_and_test.yml | 6 +++---
 qermit/taskgraph/mittask.py          | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index 3167a138..a372075c 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -87,9 +87,9 @@ jobs:
     - name: Checkout
       uses: actions/checkout@v4
     - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: 3.12
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.12
     - name: Build qermit
       run: |
         pip install .[tests]
diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 67d7531a..8cbcba71 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -16,7 +16,7 @@
 from collections import namedtuple
 from enum import Enum
 from types import MethodType
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from pytket import Bit, Circuit, Qubit
 from pytket.backends import ResultHandle

From 43d4838deec33799d63988f87d1a92e140a82fbd Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:15:26 +0000
Subject: [PATCH 8/9] Remove typevar

---
 qermit/postselection/postselect_manager.py | 3 +--
 qermit/taskgraph/mittask.py                | 5 -----
 tests/leakage_gadget_test.py               | 2 +-
 3 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/qermit/postselection/postselect_manager.py b/qermit/postselection/postselect_manager.py
index d2c9a460..2702a7b5 100644
--- a/qermit/postselection/postselect_manager.py
+++ b/qermit/postselection/postselect_manager.py
@@ -4,10 +4,9 @@
 from pytket.backends.backendresult import BackendResult
 from pytket.circuit import Bit
 from pytket.utils.outcomearray import OutcomeArray
-from qermit.taskgraph.mittask import PostselectMgrType
 
 
-class PostselectMgr[PostselectMgrType]:
+class PostselectMgr:
     """Class for tracking and applying post selection to results.
 
     An example use case might be the following. Here a Bell state is
diff --git a/qermit/taskgraph/mittask.py b/qermit/taskgraph/mittask.py
index 8cbcba71..9e862527 100644
--- a/qermit/taskgraph/mittask.py
+++ b/qermit/taskgraph/mittask.py
@@ -22,7 +22,6 @@
 from pytket.backends import ResultHandle
 from pytket.backends.backendresult import BackendResult
 from pytket.utils import QubitPauliOperator
-from typing import TypeVar
 
 
 class IOTask(Enum):
@@ -40,10 +39,7 @@ class IOTask(Enum):
     "ObservableExperiment", ["AnsatzCircuit", "ObservableTracker"]
 )
 
-PostselectMgrType = TypeVar('PostselectMgrType')
-
 Wire = Union[
-    PostselectMgrType,
     CircuitShots,
     Circuit,
     BackendResult,
@@ -60,7 +56,6 @@ class IOTask(Enum):
 ]
 
 WireList = Union[
-    List[PostselectMgrType],
     List[CircuitShots],
     List[Circuit],
     List[BackendResult],
diff --git a/tests/leakage_gadget_test.py b/tests/leakage_gadget_test.py
index 70096471..dbb346ae 100644
--- a/tests/leakage_gadget_test.py
+++ b/tests/leakage_gadget_test.py
@@ -77,7 +77,7 @@ def test_compare_with_prune() -> None:
     qermit_result_list = postselection_task(
         (
             result_list,
-            postselect_mgr_list,
+            postselect_mgr_list,  # type: ignore
         )
     )
     pytket_result_list = [

From 9fafe5ad340bb52cbfcccc7ccb74a7e9e19dadc3 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:18:47 +0000
Subject: [PATCH 9/9] Formatting

---
 .github/workflows/build_and_test.yml | 4 ----
 tests/leakage_gadget_test.py         | 8 +++++---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index a372075c..f02b7981 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -86,10 +86,6 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: 3.12
     - name: Build qermit
       run: |
         pip install .[tests]
diff --git a/tests/leakage_gadget_test.py b/tests/leakage_gadget_test.py
index dbb346ae..ff0194ad 100644
--- a/tests/leakage_gadget_test.py
+++ b/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,
@@ -26,8 +28,6 @@
 from qermit.postselection.postselect_mitres import gen_postselect_task
 from qermit.taskgraph import gen_compiled_MitRes
 
-from typing import cast
-
 
 def test_leakage_gadget() -> None:
     backend = MockQuantinuumBackend()
@@ -71,7 +71,9 @@ def test_compare_with_prune() -> None:
             ),
             n_shots=detection_circuit_shots.Shots,
         )
-        for detection_circuit_shots in cast(list[CircuitShots], detection_circuit_shots_list)
+        for detection_circuit_shots in cast(
+            list[CircuitShots], detection_circuit_shots_list
+        )
     ]
 
     qermit_result_list = postselection_task(