fix type checking

microsoft · May 24, 2024 · b6a850d · b6a850d
1 parent bbea473
commit b6a850d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/qcodes/instrument_drivers/Keithley/Keithley_2450.py b/src/qcodes/instrument_drivers/Keithley/Keithley_2450.py
@@ -698,6 +698,7 @@ def _set_sense_function(self, value: str) -> None:
         self.write(
             f":SENS:FUNC {value}",
         )
+        assert self.sense_function.inverse_val_mapping is not None
         sense_function = self.sense_function.inverse_val_mapping[value]
         sense = self.submodules[f"_sense_{sense_function}"]
         if not isinstance(sense, Keithley2450Sense):
@@ -728,8 +729,9 @@ def _set_source_function(self, value: str) -> None:
             )
 
         self.write(f":SOUR:FUNC {value}")
+        assert self.source_function.inverse_val_mapping is not None
         source_function = self.source_function.inverse_val_mapping[value]
-        source = self.submodules[f"_source_{source_function}"]
+        source = cast(Keithley2450Source, self.submodules[f"_source_{source_function}"])
         self.sense.sweep.setpoints = (source.sweep_axis,)
         if not isinstance(source, Keithley2450Source):
             raise RuntimeError(