add type hints to Experiment.py and dac.py, set up pyrefly as dev dependency

src/damaris/experiments/Experiment.py

@@ -3,20 +3,20 @@ import types
 import numpy
 
 class StateBase(object):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
-    def to_xml(self, indent = ""):
+    def to_xml(self, indent: str = "") -> str:
         return indent + "<!-- " + repr(self) + " -->"
 
 class StateSimple(StateBase):
-    def __init__(self, time, content=None):
+    def __init__(self, time: float, content: object | None = None) -> None:
         super(StateSimple, self).__init__()
         if time < 0:
             raise AssertionError("time for state is negative!")
         self.time = time
         self.content = content
 
-    def to_xml(self, indent = ""):
+    def to_xml(self, indent: str = "") -> str:
         s = indent + '<state time="%s"' % repr(self.time)
         if self.content is None:
             return s + '/>\n'
@@ -24,14 +24,14 @@ class StateSimple(StateBase):
         s += indent + '  ' + str(self.content) + '\n'
         s += indent + '</state>\n'
         return s
-    def __repr__(self):
+    def __repr__(self) -> str:
         return 'StateSimple(%s, %s)' % (self.time, repr(self.content))
 
 class StateList(StateBase):
-    def __init__(self):
+    def __init__(self) -> None:
         super(StateList, self).__init__()
-        self.list = []
-    def to_xml(self, indent = "  "):
+        self.list: list[object] = []
+    def to_xml(self, indent: str = "  ") -> str:
         s = ""
         for k in self.list:
             if hasattr(k, "to_xml"):
@@ -39,21 +39,21 @@ class StateList(StateBase):
             else:
                 s += indent + str(k)
         return s
-    def append(self, val):
+    def append(self, val: object) -> None:
         self.list.append(val)
 
 
 class StateLoop(StateList):
     """Represents a loop in the state tree"""
-    def __init__(self, repeat):
+    def __init__(self, repeat: int) -> None:
         super(StateLoop, self).__init__()
         self.repeat = repeat
-    def to_xml(self, indent = ""):
+    def to_xml(self, indent: str = "") -> str:
         s = indent + ('<sequent repeat="%d">\n' % self.repeat)
         s += super(StateLoop, self).to_xml(indent + "  ")
         s += indent + '</sequent>\n'
         return s
-    def __repr__(self):
+    def __repr__(self) -> str:
         return 'StateLoop(repeat=%d, %s)' \
             % (self.repeat, repr(self.list))
 
@@ -78,32 +78,32 @@ class Experiment:
     :param list rf_gates: list of rf gates ttl channels  len(rf_gates) == len(rf_sources)
 
     """
-    job_id = 0
+    job_id: int = 0
 
-    def __init__(self, gating=None, rf_sources=[], rf_gates=[]):
-        self.job_id = Experiment.job_id
+    def __init__(self, gating: float | None = None, rf_sources: list[int] = [], rf_gates: list[int] = []) -> None:
+        self.job_id: int = Experiment.job_id
         Experiment.job_id += 1
 
-        self.state_list = StateList()
-        self.list_stack = []
-        self.description = { }
-        self.gating = gating
+        self.state_list: StateList = StateList()
+        self.list_stack: list[StateList] = []
+        self.description: dict[str, object] = { }
+        self.gating: float | None = gating
         assert type(rf_sources) == type(list()), "rf_sources needs to be a list with the channels"
         assert type(rf_gates) == type(list()), "rf_gates needs to be a list with the channels"
         assert len(rf_gates) == len(rf_sources), "rf_sources and rf_gates must have equal number of entries"
-        self.rf_sources = rf_sources
-        self.rf_gates = rf_gates
+        self.rf_sources: list[int] = rf_sources
+        self.rf_gates: list[int] = rf_gates
         
         #for tracking the experiment length:
         #because loops are possible we need to track the length for each loop level
-        self.total_time=[0.0]
+        self.total_time: list[float] = [0.0]
         
         #and we need to know the number of iterations of the loops to multiply the state.
-        self.loop_iterations=[1]
+        self.loop_iterations: list[int] = [1]
 
     # Commands -------------------------------------------------------------------------------------
 
-    def ttl_pulse(self, length, channel = None, value = None):
+    def ttl_pulse(self, length: float, channel: int | None = None, value: int | None = None) -> None:
         """
         Creates a state with length **length** and switches requested bits of the pulse programmer to HIGH.
 
@@ -151,7 +151,7 @@ class Experiment:
         self.total_time[-1] += length
 
     ## Same as ttl_pulse, but no *channel* keyword
-    def ttls(self, length = None, value = None):
+    def ttls(self, length: float | None = None, value: int | None = None) -> None:
         """
         Same as ttl_pulse, but no *channel* keyword
 
@@ -159,6 +159,7 @@ class Experiment:
         :param int value: lines to set (integer)
         :return:
         """
+        assert value is not None, "ttls needs an integer value"
         the_value=int(value)
         s_content = '<ttlout value="0x%06x"/>' % the_value
         if length is not None:
@@ -168,7 +169,7 @@ class Experiment:
             self.state_list.append(s_content)
             
 
-    def rf_pulse(self, length=None, phase=0, source=0):
+    def rf_pulse(self, length: float | None = None, phase: float = 0, source: int = 0) -> None:
         """
         Make an rf pulse, including gating and phase switching. Only possible if Experiment is configured.
 
@@ -183,7 +184,7 @@ class Experiment:
         self.ttls(length=length, value=self.rf_gates[source]+self.rf_sources[source])
 
     ## Beginning of a new state
-    def state_start(self, time):
+    def state_start(self, time: float) -> None:
         """
         Starts a state in the pulse programs with duration *time*.
         This must be closed with :func:`state_end`.
@@ -194,13 +195,13 @@ class Experiment:
         self.total_time[-1] += time
 
     ## End of *state_start*
-    def state_end(self):
+    def state_end(self) -> None:
         """
         Closes a previous :func:`state_start`
         """
         self.state_list.append('</state>\n')
 
-    def wait(self, time, ttls=None, gating=False):
+    def wait(self, time: float, ttls: int | None = None, gating: bool = False) -> None:
         """
         Waits for a specified amount of time without performing any actions.
 
@@ -241,6 +242,7 @@ class Experiment:
 
         """
         if gating:
+            assert self.gating is not None, "gating time must be set for gating=True"
             time -= self.gating
         if ttls is not None:
             s_content = '<ttlout value="0x%06x"/>' % ttls
@@ -251,7 +253,7 @@ class Experiment:
         self.total_time[-1] += time
 
 
-    def record(self, samples, frequency, timelength=None, sensitivity = None, ttls=None, channels = 3, offset = None, impedance = None):
+    def record(self, samples: int, frequency: float, timelength: float | None = None, sensitivity: float | list[float] | None = None, ttls: int | None = None, channels: int = 3, offset: int | list[int] | None = None, impedance: float | list[float] | None = None) -> None:
         """
         Records data with a given number of samples, sampling frequency, and sensitivity.
         This command starts data acquisition from the ADC (Analog-to-Digital Converter).
@@ -365,7 +367,7 @@ class Experiment:
 
     ## Create a loop on the pulse programmer. Loop contents can not change inside the loop.
     # @params iterations Number of loop iterations
-    def loop_start(self, iterations):
+    def loop_start(self, iterations: int) -> None:
         """
         creates a loop of given number of iterations and has to be closed by loop_end().
         Commands inside the loop can not change, i.e., the parameters are the same for each loop run.
@@ -384,7 +386,7 @@ class Experiment:
         self.loop_iterations.append(iterations)
 
     ## End loop state
-    def loop_end(self):
+    def loop_end(self) -> None:
         """
         Closes a l:func:`loop_start` statement.
         """
@@ -396,7 +398,7 @@ class Experiment:
         self.total_time[-1] += looptime*loopiterations
 
 
-    def set_frequency(self, frequency, phase, ttls=0):
+    def set_frequency(self, frequency: float, phase: float, ttls: int = 0) -> None:
         """
         Sets the frequency and phase of the frequency source and optionally further channels.
         The state length is 2 µs. Stabilisation time depends on RF source.
@@ -412,7 +414,7 @@ class Experiment:
         
         self.total_time[-1] += 2e-6
 
-    def set_pfg(self, dac_value=None, length=None, shape=('rec', 0), trigger=4, is_seq=False):
+    def set_pfg(self, dac_value: int | None = None, length: float | None = None, shape: tuple[str, float] = ('rec', 0), trigger: int = 4, is_seq: bool = False) -> None:
         """
         This sets the value for the PFG, it also sets it back to dac_value=0 automatically.
         If you don't wish to do so (i.e. line shapes)  set is_seq=1
@@ -480,7 +482,7 @@ class Experiment:
         else: # don't know what to do
             raise SyntaxError("form is unknown: %s"%form)
 
-    def set_dac(self, dac_value, dac_id=1, length=None, is_seq=False, ttls=0):
+    def set_dac(self, dac_value: int, dac_id: int = 1, length: float | None = None, is_seq: bool = False, ttls: int = 0) -> None:
         """
         Sets the value for the DAC (Digital-to-Analog Converter) and optionally additional TTL lines.
         This command is used to control analog output signals, such as pulsed field gradients.
@@ -544,7 +546,7 @@ class Experiment:
             
             self.total_time[-1] += 42*9e-8
 
-    def set_phase(self, phase, ttls=0):
+    def set_phase(self, phase: float, ttls: int = 0) -> None:
         """
         Sets the phase of the RF source to the specified value.
 
@@ -589,7 +591,7 @@ class Experiment:
         self.total_time[-1] += 0.5e-6
 
 
-    def set_description(self, key, value):
+    def set_description(self, key: str, value: object) -> None:
         """
         Sets a description key-value pair in the experiment description dictionary.
 
@@ -629,7 +631,7 @@ class Experiment:
             print('Warning: Overwriting existing description "%s" = "%s" with "%s"' % (key, self.description[key], value))
         self.description[key] = value
 
-    def set_pts_local(self):
+    def set_pts_local(self) -> None:
         """
         Sets the PTS310/PTS500 frequency source to local mode.
         State length is 2 µs.
@@ -653,7 +655,7 @@ class Experiment:
         return self.job_id
 
 
-    def write_xml_string(self):
+    def write_xml_string(self) -> str:
         """
         Returns the current program as a string
 
@@ -706,14 +708,14 @@ class Experiment:
 
         return xml_string
 
-    def write_quit_job(self):
+    def write_quit_job(self) -> str:
         """
         Returns a xml quit-job
         :returns: XML quit-job
         """
         return '<?xml version="1.0" encoding="ISO-8859-1"?>\n<quit/>'
     
-    def get_length(self):
+    def get_length(self) -> float:
         timelength = 0.0
         #calculate the correct timelength also for unclosed loops, if there is no unclosed loop
         #the timelength of this experiment is self.total_time[-1].
@@ -723,13 +725,13 @@ class Experiment:
         return timelength
 
 class Quit(Experiment):
-    def write_xml_string(self):
+    def write_xml_string(self) -> str:
         return '<?xml version="1.0" encoding="ISO-8859-1"?>\n<quit no="%d"/>'%self.job_id
 
 
 # /Public Methods ------------------------------------------------------------------------------
 
-def self_test():
+def self_test() -> None:
     e = Experiment()
     e.set_description("key", "value")
     e.set_frequency(85e6, 90, ttls=16)