Fix lakeshore336/335/372 snapshot bug

src/qcodes/instrument/sims/lakeshore_model336.yaml

@@ -301,20 +301,20 @@ devices:
           r: "{}"
 
       htrset_output_1:
-        default: "1, 5"
+        default: "2,2,+1.7320E+00,1"
         getter:
           q: "HTRSET? 1"
           r: "{}"
         setter:
-          q: "HTRSET 1, {}"
+          q: "HTRSET 1,{}"
 
       htrset_output_2:
-        default: "1, 5"
+        default: "2,2,+1.7320E+00,1"
         getter:
           q: "HTRSET? 2"
           r: "{}"
         setter:
-          q: "HTRSET 2, {}"
+          q: "HTRSET 2,{}"
 
       ramp_output_1:
         default: "0,0"

src/qcodes/instrument/sims/lakeshore_model372.yaml

@@ -898,6 +898,14 @@ devices:
         setter:
           q: "SETP 0,{}"
 
+      htrset_output_0:
+        default: "+120.000,0,+0.000,2"
+        getter:
+          q: "HTRSET? 0"
+          r: "{}"
+        setter:
+          q: "HTRSET 0,{}"
+
       # ====================
       # Heater Output 1 (warmup_heater)
       # ====================
@@ -933,6 +941,14 @@ devices:
         setter:
           q: "SETP 1,{}"
 
+      htrset_output_1:
+        default: "2,1,+0.000,1"
+        getter:
+          q: "HTRSET? 1"
+          r: "{}"
+        setter:
+          q: "HTRSET 1,{}"
+
       # ====================
       # Heater Output 2 (analog_heater)
       # ====================

src/qcodes/instrument_drivers/Lakeshore/Lakeshore_model_336.py

@@ -6,6 +6,7 @@
 from .lakeshore_base import (
     LakeshoreBase,
     LakeshoreBaseOutput,
+    LakeshoreBaseOutputWithHeaterSetup,
     LakeshoreBaseSensorChannel,
 )
 
@@ -31,12 +32,13 @@
 }
 
 
-class LakeshoreModel336CurrentSource(LakeshoreBaseOutput):
+class LakeshoreModel336CurrentSource(LakeshoreBaseOutputWithHeaterSetup):
     """
     InstrumentChannel for current sources on Lakeshore Model 336.
 
     Class for control outputs 1 and 2 of Lakeshore Model 336 that are variable DC current
-    sources referenced to chassis ground.
+    sources referenced to chassis ground. Unlike the Model 335, the heater setup
+    (``HTRSET``) command of the Model 336 has no output type field.
     """
 
     MODES: ClassVar[dict[str, int]] = {
@@ -71,7 +73,7 @@ class LakeshoreModel336VoltageSource(LakeshoreBaseOutput):
     InstrumentChannel for voltage sources on Lakeshore Model 336.
 
     This is used for control outputs 3 and 4 that are variable DC voltage
-    sources.
+    sources. These outputs do not support the heater setup (``HTRSET``) command.
     """
 
     MODES: ClassVar[dict[str, int]] = {

src/qcodes/instrument_drivers/Lakeshore/Lakeshore_model_372.py

@@ -18,7 +18,15 @@
 
 
 class LakeshoreModel372Output(LakeshoreBaseOutput):
-    """An InstrumentChannel for control outputs (heaters) of Lakeshore Model 372"""
+    """
+    An InstrumentChannel for control outputs (heaters) of Lakeshore Model 372
+
+    This class is used for the analog (still) heater, which does not support
+    the heater setup (``HTRSET``) command. The sample heater and the warm-up
+    heater use the subclasses :class:`LakeshoreModel372SampleHeater` and
+    :class:`LakeshoreModel372WarmupHeater`, which add the parameters of that
+    command.
+    """
 
     MODES: ClassVar[dict[str, int]] = {
         "off": 0,
@@ -105,6 +113,121 @@ def __init__(
         self.D.vals = vals.Numbers(0, 2500)
 
 
+class LakeshoreModel372SampleHeater(LakeshoreModel372Output):
+    """
+    An InstrumentChannel for the sample heater (output 0) of Lakeshore Model 372
+
+    Unlike the Model 335, the heater setup (``HTRSET``) command of the
+    Model 372 has no output type field. The max current fields of this command
+    only apply to the warm-up heater, hence the sample heater only has the
+    heater resistance and current/power display parameters.
+    """
+
+    def __init__(
+        self,
+        parent: "LakeshoreModel372",
+        output_name: str,
+        output_index: int,
+        **kwargs: "Unpack[InstrumentBaseKWArgs]",
+    ) -> None:
+        super().__init__(parent, output_name, output_index, **kwargs)
+
+        self.output_heater_resistance: GroupParameter = self.add_parameter(
+            name="output_heater_resistance",
+            docstring="Heater load in ohms (1 to 2000)",
+            vals=vals.Numbers(1, 2000),
+            unit="Ohm",
+            get_parser=float,
+            parameter_class=GroupParameter,
+        )
+        """Heater load in ohms (1 to 2000)"""
+
+        self.output_display: GroupParameter = self.add_parameter(
+            name="output_display",
+            docstring="Specifies whether the heater output displays in current or power",
+            val_mapping={"current": 1, "power": 2},
+            parameter_class=GroupParameter,
+        )
+        """Specifies whether the heater output displays in current or power"""
+
+        # The manual requires zeros for the max current and max user current
+        # fields when configuring the sample heater
+        self.heater_group = Group(
+            [self.output_heater_resistance, self.output_display],
+            set_cmd=f"HTRSET {output_index},{{output_heater_resistance}},0,0,{{output_display}}",
+            get_cmd=f"HTRSET? {output_index}",
+            get_parser=self._parse_heater_setup,
+        )
+
+    @staticmethod
+    def _parse_heater_setup(response: str) -> dict[str, str]:
+        resistance, _, _, display = response.split(",")
+        return {"output_heater_resistance": resistance, "output_display": display}
+
+
+class LakeshoreModel372WarmupHeater(LakeshoreModel372Output):
+    """
+    An InstrumentChannel for the warm-up heater (output 1) of Lakeshore Model 372
+
+    Unlike the Model 335, the heater setup (``HTRSET``) command of the
+    Model 372 has no output type field.
+    """
+
+    def __init__(
+        self,
+        parent: "LakeshoreModel372",
+        output_name: str,
+        output_index: int,
+        **kwargs: "Unpack[InstrumentBaseKWArgs]",
+    ) -> None:
+        super().__init__(parent, output_name, output_index, **kwargs)
+
+        self.output_heater_resistance: GroupParameter = self.add_parameter(
+            name="output_heater_resistance",
+            docstring="Heater Resistance Setting: 25/50ohm",
+            val_mapping={"25ohm": 1, "50ohm": 2},
+            parameter_class=GroupParameter,
+        )
+        """Heater Resistance Setting: 25/50ohm"""
+
+        self.output_max_current: GroupParameter = self.add_parameter(
+            name="output_max_current",
+            docstring="Specifies the maximum heater output current: User Specified, 0.45 A, 0.63 A",
+            val_mapping={"user": 0, "0.45A": 1, "0.63A": 2},
+            parameter_class=GroupParameter,
+        )
+        """Specifies the maximum heater output current: User Specified, 0.45 A, 0.63 A"""
+
+        self.output_max_user_current: GroupParameter = self.add_parameter(
+            name="output_max_user_current",
+            docstring="Specifies the maximum heater output current if max current is set to User Specified.",
+            vals=vals.Numbers(0, 0.63),
+            unit="A",
+            get_parser=float,
+            parameter_class=GroupParameter,
+        )
+        """Specifies the maximum heater output current if max current is set to User Specified."""
+
+        self.output_display: GroupParameter = self.add_parameter(
+            name="output_display",
+            docstring="Specifies whether the heater output displays in current or power",
+            val_mapping={"current": 1, "power": 2},
+            parameter_class=GroupParameter,
+        )
+        """Specifies whether the heater output displays in current or power"""
+
+        self.heater_group = Group(
+            [
+                self.output_heater_resistance,
+                self.output_max_current,
+                self.output_max_user_current,
+                self.output_display,
+            ],
+            set_cmd=f"HTRSET {output_index},{{output_heater_resistance}},{{output_max_current}},{{output_max_user_current}},{{output_display}}",
+            get_cmd=f"HTRSET? {output_index}",
+        )
+
+
 class LakeshoreModel372Channel(LakeshoreBaseSensorChannel):
     """
     An InstrumentChannel representing a single sensor on a Lakeshore Model 372.
@@ -351,16 +474,20 @@ def __init__(
 
         heaters = {"sample_heater": 0, "warmup_heater": 1, "analog_heater": 2}
 
-        self.sample_heater: LakeshoreModel372Output = self.add_submodule(
+        self.sample_heater: LakeshoreModel372SampleHeater = self.add_submodule(
             "sample_heater",
-            LakeshoreModel372Output(self, "sample_heater", heaters["sample_heater"]),
+            LakeshoreModel372SampleHeater(
+                self, "sample_heater", heaters["sample_heater"]
+            ),
         )
         """
         Sample heater output channel.
         """
-        self.warmup_heater: LakeshoreModel372Output = self.add_submodule(
+        self.warmup_heater: LakeshoreModel372WarmupHeater = self.add_submodule(
             "warmup_heater",
-            LakeshoreModel372Output(self, "warmup_heater", heaters["warmup_heater"]),
+            LakeshoreModel372WarmupHeater(
+                self, "warmup_heater", heaters["warmup_heater"]
+            ),
         )
         """
         Warm-up heater output channel.

src/qcodes/instrument_drivers/Lakeshore/__init__.py

@@ -6,6 +6,8 @@
 from .lakeshore_base import (
     LakeshoreBase,
     LakeshoreBaseOutput,
+    LakeshoreBaseOutputWithHeaterSetup,
+    LakeshoreBaseOutputWithHeaterSetupAndOutputType,
     LakeshoreBaseSensorChannel,
 )
 from .Lakeshore_model_325 import (
@@ -25,11 +27,15 @@
     LakeshoreModel372,
     LakeshoreModel372Channel,
     LakeshoreModel372Output,
+    LakeshoreModel372SampleHeater,
+    LakeshoreModel372WarmupHeater,
 )
 
 __all__ = [
     "LakeshoreBase",
     "LakeshoreBaseOutput",
+    "LakeshoreBaseOutputWithHeaterSetup",
+    "LakeshoreBaseOutputWithHeaterSetupAndOutputType",
     "LakeshoreBaseSensorChannel",
     "LakeshoreModel325",
     "LakeshoreModel325Curve",
@@ -46,4 +52,6 @@
     "LakeshoreModel372",
     "LakeshoreModel372Channel",
     "LakeshoreModel372Output",
+    "LakeshoreModel372SampleHeater",
+    "LakeshoreModel372WarmupHeater",
 ]

src/qcodes/instrument_drivers/Lakeshore/_lakeshore_model_335.py

@@ -8,7 +8,7 @@
 
 from .lakeshore_base import (
     LakeshoreBase,
-    LakeshoreBaseOutput,
+    LakeshoreBaseOutputWithHeaterSetupAndOutputType,
     LakeshoreBaseSensorChannel,
 )
 
@@ -147,7 +147,7 @@ def __init__(
         )
 
 
-class LakeshoreModel335CurrentSource(LakeshoreBaseOutput):
+class LakeshoreModel335CurrentSource(LakeshoreBaseOutputWithHeaterSetupAndOutputType):
     """
     InstrumentChannel for current sources on Lakeshore Model 335.