add porots containers tests

quam/components/ports/ports_containers.py

@@ -1,5 +1,6 @@
-from typing import ClassVar, Dict, Union
+from typing import ClassVar, Dict, Optional, TypeVar, Union
 from dataclasses import field
+from quam.components.ports.base_ports import FEMPort
 from quam.core import quam_dataclass, QuamComponent
 from .analog_outputs import (
     OPXPlusAnalogOutputPort,
@@ -17,6 +18,20 @@
 
 __all__ = ["OPXPlusPortsContainer", "FEMPortsContainer"]
 
+OPXPlusPortTypes = Union[
+    OPXPlusAnalogInputPort,
+    OPXPlusAnalogOutputPort,
+    OPXPlusDigitalOutputPort,
+    OPXPlusDigitalInputPort,
+]
+FEMPortTypes = Union[
+    LFFEMAnalogInputPort,
+    LFFEMAnalogOutputPort,
+    MWFEMAnalogInputPort,
+    MWFEMAnalogOutputPort,
+    FEMDigitalOutputPort,
+]
+
 
 @quam_dataclass
 class OPXPlusPortsContainer(QuamComponent):
@@ -48,10 +63,12 @@ def _get_port(
         except KeyError:
             if not create:
                 raise KeyError(
-                    f"Port {port_id} not found in controller {controller_id}"
+                    f"Could not find existing {port_type} port: "
+                    f"{port_type} ({controller_id}, {port_id}"
                 )
 
-        ports = controllers.setdefault(controller_id, {})
+        controllers.setdefault(controller_id, {})
+        ports = controllers[controller_id]
 
         if port_type == "analog_output":
             ports[port_id] = OPXPlusAnalogOutputPort(controller_id, port_id, **kwargs)
@@ -66,6 +83,19 @@ def _get_port(
 
         return ports[port_id]
 
+    def get_port_from_reference(
+        self, port_reference: str, create=False
+    ) -> OPXPlusPortTypes:
+        elems = port_reference.split("/")
+        port_type, controller_id, port_id = elems[-3:]
+
+        if controller_id.isdigit():
+            controller_id = int(controller_id)
+        fem_id = int(fem_id)
+        port_id = int(port_id)
+
+        return self._get_port(controller_id, fem_id, port_id, port_type, create=create)
+
     def get_analog_output(
         self,
         controller_id: Union[str, int],
@@ -141,19 +171,33 @@ def _get_port(
     ):
         controllers = getattr(self, f"{port_type}s")
 
-        if not create:
+        try:
             return controllers[controller_id][fem_id][port_id]
+        except KeyError:
+            if not create:
+                raise KeyError(
+                    f"Could not find existing {port_type} port: "
+                    f"{port_type} ({controller_id}, {fem_id}, {port_id}"
+                )
 
-        fems = controllers.setdefault(controller_id, {})
-        ports = fems.setdefault(controller_id, {})
+        controllers.setdefault(controller_id, {})
+        fems = controllers[controller_id]
+        fems.setdefault(fem_id, {})
+        ports = fems[fem_id]
 
         if port_type == "analog_output":
-            ports[port_id] = LFFEMAnalogOutputPort(controller_id, fem_id, port_id)
+            ports[port_id] = LFFEMAnalogOutputPort(
+                controller_id, fem_id, port_id, **kwargs
+            )
         elif port_type == "analog_input":
-            ports[port_id] = LFFEMAnalogInputPort(controller_id, fem_id, port_id)
+            ports[port_id] = LFFEMAnalogInputPort(
+                controller_id, fem_id, port_id, **kwargs
+            )
         elif port_type == "mw_output":
+            if "upconverter_frequency" not in kwargs and "upconverters" not in kwargs:
+                kwargs["upconverter_frequency"] = 5e9
             ports[port_id] = MWFEMAnalogOutputPort(
-                controller_id, fem_id, port_id, band=kwargs.get("band", 1)
+                controller_id, fem_id, port_id, band=kwargs.get("band", 1), **kwargs
             )
         elif port_type == "mw_input":
             ports[port_id] = MWFEMAnalogInputPort(
@@ -162,14 +206,39 @@ def _get_port(
                 port_id,
                 band=kwargs.get("band", 1),  # TODO Are default values the best here?
                 downconverter_frequency=kwargs.get("downconverter_frequency", 5e9),
+                **kwargs,
             )
         elif port_type == "digital_output":
-            ports[port_id] = FEMDigitalOutputPort(controller_id, fem_id, port_id)
+            ports[port_id] = FEMDigitalOutputPort(
+                controller_id, fem_id, port_id, **kwargs
+            )
         else:
             raise ValueError(f"Invalid port type: {port_type}")
 
         return ports[port_id]
 
+    def get_port_from_reference(
+        self,
+        port_reference: Union[QuamComponent, str],
+        attr: Optional[str] = None,
+        create=False,
+    ) -> FEMPortTypes:
+        if isinstance(port_reference, QuamComponent):
+            reference = port_reference.get_reference(attr=attr)
+            if reference is None:
+                raise ValueError("Cannot get port from reference {port_reference}")
+            port_reference = reference
+
+        elems = port_reference.split("/")
+        port_type, controller_id, fem_id, port_id = elems[-4:]
+
+        if controller_id.isdigit():
+            controller_id = int(controller_id)
+        fem_id = int(fem_id)
+        port_id = int(port_id)
+
+        return self._get_port(controller_id, fem_id, port_id, port_type, create=create)
+
     def get_analog_output(
         self,
         controller_id: Union[str, int],

tests/components/ports/test_ports_containers.py

@@ -0,0 +1,106 @@
+import pytest
+from quam.components.ports import (
+    OPXPlusPortsContainer,
+    FEMPortsContainer,
+    OPXPlusAnalogOutputPort,
+    LFFEMAnalogOutputPort,
+)
+from quam.components.ports.analog_inputs import (
+    LFFEMAnalogInputPort,
+    MWFEMAnalogInputPort,
+    OPXPlusAnalogInputPort,
+)
+from quam.components.ports.analog_outputs import MWFEMAnalogOutputPort
+from quam.components.ports.digital_inputs import OPXPlusDigitalInputPort
+from quam.components.ports.digital_outputs import (
+    FEMDigitalOutputPort,
+    OPXPlusDigitalOutputPort,
+)
+
+
+def test_opx_plus_ports_container_initialize():
+    ports_container = OPXPlusPortsContainer()
+    assert ports_container.analog_outputs == {}
+    assert ports_container.analog_inputs == {}
+    assert ports_container.digital_outputs == {}
+    assert ports_container.digital_inputs == {}
+
+
+def test_fem_ports_container_initialize():
+    ports_container = FEMPortsContainer()
+    assert ports_container.analog_outputs == {}
+    assert ports_container.analog_inputs == {}
+    assert ports_container.mw_outputs == {}
+    assert ports_container.mw_inputs == {}
+    assert ports_container.digital_outputs == {}
+
+
+@pytest.mark.parametrize(
+    "port_type", ["analog_output", "analog_input", "digital_output", "digital_input"]
+)
+def test_opx_plus_ports_container_add_ports(port_type):
+    port_mapping = {
+        "analog_output": OPXPlusAnalogOutputPort,
+        "analog_input": OPXPlusAnalogInputPort,
+        "digital_output": OPXPlusDigitalOutputPort,
+        "digital_input": OPXPlusDigitalInputPort,
+    }
+
+    ports_container = OPXPlusPortsContainer()
+
+    get_port_func = getattr(ports_container, f"get_{port_type}")
+
+    with pytest.raises(KeyError):
+        get_port_func("con1", 2)
+
+    port = get_port_func("con2", 3, create=True)
+    assert isinstance(port, port_mapping[port_type])
+
+    assert port.controller_id == "con2"
+    assert port.port_id == 3
+
+    port2 = get_port_func("con2", 3, create=False)
+    assert port is port2
+
+    port3 = get_port_func("con2", 3, create=True)
+    assert port is port3
+
+    ports_group = getattr(ports_container, f"{port_type}s")
+    assert port is ports_group["con2"][3]
+
+
+@pytest.mark.parametrize(
+    "port_type",
+    ["analog_output", "analog_input", "mw_output", "mw_input", "digital_output"],
+)
+def test_fem_ports_container_add_ports(port_type):
+    port_mapping = {
+        "analog_output": LFFEMAnalogOutputPort,
+        "analog_input": LFFEMAnalogInputPort,
+        "mw_output": MWFEMAnalogOutputPort,
+        "mw_input": MWFEMAnalogInputPort,
+        "digital_output": FEMDigitalOutputPort,
+    }
+
+    ports_container = FEMPortsContainer()
+
+    get_port_func = getattr(ports_container, f"get_{port_type}")
+
+    with pytest.raises(KeyError):
+        get_port_func("con1", 1, 2)
+
+    port = get_port_func("con2", 2, 3, create=True)
+    assert isinstance(port, port_mapping[port_type])
+
+    assert port.controller_id == "con2"
+    assert port.fem_id == 2
+    assert port.port_id == 3
+
+    port2 = get_port_func("con2", 2, 3, create=False)
+    assert port is port2
+
+    port3 = get_port_func("con2", 2, 3, create=True)
+    assert port is port3
+
+    ports_group = getattr(ports_container, f"{port_type}s")
+    assert port is ports_group["con2"][2][3]