From 4d00aebb3855dcb8b0dfca9f5b7050c062d2ee2f Mon Sep 17 00:00:00 2001
From: Stavros Efthymiou <35475381+stavros11@users.noreply.github.com>
Date: Thu, 17 Oct 2024 17:59:26 +0400
Subject: [PATCH] chore: use TypedDict for type hints

---
 .../_core/instruments/qm/config/elements.py   | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/qibolab/_core/instruments/qm/config/elements.py b/src/qibolab/_core/instruments/qm/config/elements.py
index 2c15b1473..f3820da2e 100644
--- a/src/qibolab/_core/instruments/qm/config/elements.py
+++ b/src/qibolab/_core/instruments/qm/config/elements.py
@@ -1,15 +1,16 @@
 from dataclasses import dataclass, field
-from typing import Literal, Union
+from typing import Union
+
+from typing_extensions import TypedDict
 
 from qibolab._core.components import Channel
 
 __all__ = ["DcElement", "RfOctaveElement", "AcquireOctaveElement", "Element"]
 
 
-StandardInOutType = tuple[str, int]
-FemInOutType = tuple[str, int, int]
-InOutType = Union[StandardInOutType, FemInOutType]
-Port = Literal["port"]
+InOutType = Union[tuple[str, int], tuple[str, int, int]]
+OctavePort = TypedDict("OpxPlusPort", {"port": tuple[str, int]})
+Port = TypedDict("Port", {"port": InOutType})
 ConnectivityType = Union[str, tuple[str, int]]
 
 
@@ -26,7 +27,7 @@ class OutputSwitch:
     """
 
 
-def _to_port(channel: Channel) -> dict[Port, InOutType]:
+def _to_port(channel: Channel) -> Port:
     """Convert a channel to the port dictionary required for the QUA config.
 
     The following syntax is assumed for ``channel.device``:
@@ -52,7 +53,7 @@ def output_switch(connectivity: ConnectivityType, port: int):
 
 @dataclass
 class DcElement:
-    singleInput: dict[Port, InOutType]
+    singleInput: Port
     intermediate_frequency: int = 0
     operations: dict[str, str] = field(default_factory=dict)
 
@@ -61,10 +62,13 @@ def from_channel(cls, channel: Channel):
         return cls(_to_port(channel))
 
 
+DigitalInputs = TypedDict("digitalInputs", {"output_switch": OutputSwitch})
+
+
 @dataclass
 class RfOctaveElement:
-    RF_inputs: dict[Port, StandardInOutType]
-    digitalInputs: dict[Literal["output_switch"], OutputSwitch]
+    RF_inputs: OctavePort
+    digitalInputs: DigitalInputs
     intermediate_frequency: int
     operations: dict[str, str] = field(default_factory=dict)
 
@@ -84,9 +88,9 @@ def from_channel(
 
 @dataclass
 class AcquireOctaveElement:
-    RF_inputs: dict[Port, StandardInOutType]
-    RF_outputs: dict[Port, StandardInOutType]
-    digitalInputs: dict[Literal["output_switch"], OutputSwitch]
+    RF_inputs: OctavePort
+    RF_outputs: OctavePort
+    digitalInputs: DigitalInputs
     intermediate_frequency: int
     time_of_flight: int = 24
     smearing: int = 0