Merge pull request #7482 from jenshnielsen/b1500_types

Improve types in b1500

Author: jenshnielsen

src/qcodes/instrument_drivers/Keysight/keysightb1500/KeysightB1520A.py

@@ -1260,7 +1260,7 @@ def __init__(self, parent: "KeysightB1520A", name: str, **kwargs: Any):
         super().__init__(parent=parent, name=name, **kwargs)
         self._chnum = parent.channels[0]
 
-        self.add_submodule(
+        self.frequency_list = self.add_submodule(
             "frequency_list",
             KeysightB1500FrequencyList(self, "frequency_list", self._chnum),
         )

tests/drivers/keysight_b1500/b1500_driver_tests/conftest.py

@@ -1,3 +1,4 @@
+from typing import TYPE_CHECKING
 from unittest.mock import MagicMock, PropertyMock
 
 import pytest
@@ -8,9 +9,12 @@
     KeysightB1500,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 
 @pytest.fixture(name="b1500")
-def _make_b1500(request: FixtureRequest):
+def _make_b1500(request: FixtureRequest) -> "Generator[KeysightB1500, None, None]":
     request.addfinalizer(KeysightB1500.close_all)
 
     try:
@@ -30,7 +34,7 @@ def _make_b1500(request: FixtureRequest):
 
 
 @pytest.fixture(name="mainframe")
-def _make_mainframe():
+def _make_mainframe() -> "Generator[MagicMock, None, None]":
     PropertyMock()
     mainframe = MagicMock()
     name_parts = PropertyMock(return_value=["mainframe"])

tests/drivers/keysight_b1500/b1500_driver_tests/test_b1500.py

@@ -26,7 +26,7 @@
 )
 
 
-def test_make_module_from_model_name(mainframe) -> None:
+def test_make_module_from_model_name(mainframe: MagicMock) -> None:
     with pytest.raises(NotImplementedError):
         KeysightB1500.from_model_name(
             model="unsupported_module", slot_nr=0, parent=mainframe, name="dummy"
@@ -53,43 +53,50 @@ def test_make_module_from_model_name(mainframe) -> None:
     assert isinstance(smu, KeysightB1511B)
 
 
-def test_init(b1500) -> None:
+def test_init(b1500: KeysightB1500) -> None:
     assert hasattr(b1500, "smu1")
     assert hasattr(b1500, "smu2")
     assert hasattr(b1500, "cmu1")
     assert hasattr(b1500, "wgfmu1")
 
 
-def test_snapshot_does_not_raise_warnings(b1500) -> None:
+def test_snapshot_does_not_raise_warnings(b1500: KeysightB1500) -> None:
     with warnings.catch_warnings():
         warnings.simplefilter("error")
         b1500.snapshot(update=True)
 
 
-def test_submodule_access_by_class(b1500) -> None:
-    assert b1500.smu1 in b1500.by_kind["SMU"]
-    assert b1500.smu2 in b1500.by_kind["SMU"]
-    assert b1500.cmu1 in b1500.by_kind["CMU"]
-    assert b1500.wgfmu1 in b1500.by_kind["WGFMU"]
+def test_submodule_access_by_class(b1500: KeysightB1500) -> None:
+    assert b1500.smu1 in b1500.by_kind[constants.ModuleKind.SMU]
+    # while it does not type check it is possible to look up by string
+    assert b1500.smu1 in b1500.by_kind["SMU"]  # type: ignore
+    assert b1500.smu1 in b1500.by_kind[constants.ModuleKind.SMU]
+    assert b1500.smu2 in b1500.by_kind["SMU"]  # type: ignore
+    assert b1500.cmu1 in b1500.by_kind[constants.ModuleKind.CMU]
+    assert b1500.cmu1 in b1500.by_kind["CMU"]  # type: ignore
+    assert b1500.wgfmu1 in b1500.by_kind[constants.ModuleKind.WGFMU]
+    assert b1500.wgfmu1 in b1500.by_kind["WGFMU"]  # type: ignore
 
 
-def test_submodule_access_by_slot(b1500) -> None:
+def test_submodule_access_by_slot(b1500: KeysightB1500) -> None:
     assert b1500.smu1 is b1500.by_slot[SlotNr.SLOT01]
     assert b1500.smu2 is b1500.by_slot[SlotNr.SLOT02]
-    assert b1500.cmu1 is b1500.by_slot[3]
-    assert b1500.wgfmu1 is b1500.by_slot[6]
+    # while it does not type check it is possible to look up by integer
+    assert b1500.cmu1 is b1500.by_slot[3]  # type: ignore
+    assert b1500.wgfmu1 is b1500.by_slot[6]  # type: ignore
 
 
-def test_submodule_access_by_channel(b1500) -> None:
+def test_submodule_access_by_channel(b1500: KeysightB1500) -> None:
     assert b1500.smu1 is b1500.by_channel[ChNr.SLOT_01_CH1]
     assert b1500.smu2 is b1500.by_channel[ChNr.SLOT_02_CH1]
     assert b1500.cmu1 is b1500.by_channel[ChNr.SLOT_03_CH1]
     assert b1500.wgfmu1 is b1500.by_channel[ChNr.SLOT_06_CH1]
-    assert b1500.wgfmu1 is b1500.by_channel[6]
+    # while it does not type check it is possible to look up by integer
+    assert b1500.wgfmu1 is b1500.by_channel[6]  # type: ignore
     assert b1500.wgfmu1 is b1500.by_channel[ChNr.SLOT_06_CH2]
 
 
-def test_enable_multiple_channels(b1500) -> None:
+def test_enable_multiple_channels(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -98,7 +105,7 @@ def test_enable_multiple_channels(b1500) -> None:
     mock_write.assert_called_once_with("CN 1,2,3")
 
 
-def test_disable_multiple_channels(b1500) -> None:
+def test_disable_multiple_channels(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -107,7 +114,7 @@ def test_disable_multiple_channels(b1500) -> None:
     mock_write.assert_called_once_with("CL 1,2,3")
 
 
-def test_use_nplc_for_high_speed_adc(b1500) -> None:
+def test_use_nplc_for_high_speed_adc(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -120,7 +127,7 @@ def test_use_nplc_for_high_speed_adc(b1500) -> None:
     mock_write.assert_called_once_with("AIT 0,2,3")
 
 
-def test_use_nplc_for_high_resolution_adc(b1500) -> None:
+def test_use_nplc_for_high_resolution_adc(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -133,7 +140,7 @@ def test_use_nplc_for_high_resolution_adc(b1500) -> None:
     mock_write.assert_called_once_with("AIT 1,2,8")
 
 
-def test_autozero_enabled(b1500) -> None:
+def test_autozero_enabled(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -150,7 +157,7 @@ def test_autozero_enabled(b1500) -> None:
     assert b1500.autozero_enabled() is False
 
 
-def test_use_manual_mode_for_high_speed_adc(b1500) -> None:
+def test_use_manual_mode_for_high_speed_adc(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -168,7 +175,7 @@ def test_use_manual_mode_for_high_speed_adc(b1500) -> None:
     mock_write.assert_called_once_with("AIT 0,1,8")
 
 
-def test_self_calibration_successful(b1500) -> None:
+def test_self_calibration_successful(b1500: KeysightB1500) -> None:
     mock_ask = MagicMock()
     b1500.ask = mock_ask
 
@@ -180,7 +187,7 @@ def test_self_calibration_successful(b1500) -> None:
     mock_ask.assert_called_once_with("*CAL?")
 
 
-def test_self_calibration_failed(b1500) -> None:
+def test_self_calibration_failed(b1500: KeysightB1500) -> None:
     mock_ask = MagicMock()
     b1500.ask = mock_ask
 
@@ -193,12 +200,12 @@ def test_self_calibration_failed(b1500) -> None:
     mock_ask.assert_called_once_with("*CAL?")
 
 
-def test_error_message(b1500) -> None:
+def test_error_message(b1500: KeysightB1500) -> None:
     response = b1500.error_message()
     assert '+0,"No Error."' == response
 
 
-def test_clear_timer_count(b1500) -> None:
+def test_clear_timer_count(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -211,15 +218,15 @@ def test_clear_timer_count(b1500) -> None:
     mock_write.assert_called_once_with("TSR 1")
 
 
-def test_set_measuremet_mode(b1500) -> None:
+def test_set_measuremet_mode(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
     b1500.set_measurement_mode(mode=constants.MM.Mode.SPOT, channels=[1, 2])
     mock_write.assert_called_once_with("MM 1,1,2")
 
 
-def test_get_measurement_mode(b1500) -> None:
+def test_get_measurement_mode(b1500: KeysightB1500) -> None:
     mock_ask = MagicMock()
     b1500.ask = mock_ask
 
@@ -229,7 +236,7 @@ def test_get_measurement_mode(b1500) -> None:
     assert measurement_mode["channels"] == [1, 2]
 
 
-def test_get_response_format_and_mode(b1500) -> None:
+def test_get_response_format_and_mode(b1500: KeysightB1500) -> None:
     mock_ask = MagicMock()
     b1500.ask = mock_ask
 
@@ -239,7 +246,7 @@ def test_get_response_format_and_mode(b1500) -> None:
     assert measurement_mode["mode"] == constants.FMT.Mode(1)
 
 
-def test_enable_smu_filters(b1500) -> None:
+def test_enable_smu_filters(b1500: KeysightB1500) -> None:
     mock_write = MagicMock()
     b1500.write = mock_write
 
@@ -271,7 +278,9 @@ def test_enable_smu_filters(b1500) -> None:
     mock_write.assert_called_once_with("FL 1,102,202,302")
 
 
-def test_error_message_is_called_after_setting_a_parameter(b1500) -> None:
+def test_error_message_is_called_after_setting_a_parameter(
+    b1500: KeysightB1500,
+) -> None:
     mock_ask = MagicMock()
     b1500.ask = mock_ask
     mock_ask.return_value = '+0,"No Error."'

tests/drivers/keysight_b1500/b1500_driver_tests/test_b1511b_smu.py

@@ -1,4 +1,5 @@
 import re
+from typing import TYPE_CHECKING
 
 import pytest
 
@@ -10,15 +11,21 @@
     KeysightB1511B,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from unittest.mock import MagicMock
+
 
 @pytest.fixture(name="smu")
-def _make_smu(mainframe):
+def _make_smu(mainframe: "MagicMock") -> "Generator[KeysightB1511B, None, None]":
     slot_nr = 1
     smu = KeysightB1511B(parent=mainframe, name="B1511B", slot_nr=slot_nr)
     yield smu
 
 
-def test_force_invalid_current_output_range_when_asu_not_present(smu) -> None:
+def test_force_invalid_current_output_range_when_asu_not_present(
+    smu: KeysightB1511B,
+) -> None:
     msg = re.escape("Invalid Source Current Output Range")
     with pytest.raises(RuntimeError, match=msg):
         smu.asu_present = True
@@ -27,7 +34,7 @@ def test_force_invalid_current_output_range_when_asu_not_present(smu) -> None:
 
 
 def test_i_measure_range_config_raises_invalid_range_error_when_asu_not_present(
-    smu,
+    smu: KeysightB1511B,
 ) -> None:
     msg = re.escape("8 current measurement range")
     with pytest.raises(RuntimeError, match=msg):