[unitaryHACK] Add type hints to _seq_drawer.py (#16)

pulser/_seq_drawer.py

@@ -12,14 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import matplotlib.pyplot as plt
 import numpy as np
+
+import pulser
 from pulser.waveforms import ConstantWaveform
 from pulser.pulse import Pulse
 from scipy.interpolate import CubicSpline
+from typing import Any, cast, Dict, Optional, Tuple, Union
 
 
-def gather_data(seq):
+def gather_data(seq: pulser.sequence.Sequence) -> Dict:
     """Collects the whole sequence data for plotting.
 
     Args:
@@ -36,7 +41,7 @@ def gather_data(seq):
         time = [-1]     # To not break the "time[-1]" later on
         amp = []
         detuning = []
-        target = {}
+        target: Dict[Union[str, Tuple[int, int]], Any] = {}
         # phase_shift = {}
         for slot in sch:
             if slot.ti == -1:
@@ -52,12 +57,12 @@ def gather_data(seq):
                 if slot.type == 'target':
                     target[(slot.ti, slot.tf-1)] = slot.targets
                 continue
-            pulse = slot.type
+            pulse = cast(Pulse, slot.type)
             if (isinstance(pulse.amplitude, ConstantWaveform) and
                     isinstance(pulse.detuning, ConstantWaveform)):
                 time += [slot.ti, slot.tf-1]
-                amp += [pulse.amplitude._value] * 2
-                detuning += [pulse.detuning._value] * 2
+                amp += [int(pulse.amplitude._value)] * 2
+                detuning += [int(pulse.detuning._value)] * 2
             else:
                 time += list(range(slot.ti, slot.tf))
                 amp += pulse.amplitude.samples.tolist()
@@ -75,7 +80,9 @@ def gather_data(seq):
     return data
 
 
-def draw_sequence(seq, sampling_rate=None, draw_phase_area=False):
+def draw_sequence(seq: pulser.sequence.Sequence,
+                  sampling_rate: Optional[float] = None,
+                  draw_phase_area: bool = False) -> None:
     """Draw the entire sequence.
 
     Args:
@@ -89,7 +96,7 @@ def draw_sequence(seq, sampling_rate=None, draw_phase_area=False):
             as text on the plot, defaults to False.
     """
 
-    def phase_str(phi):
+    def phase_str(phi: float) -> str:
         """Formats a phase value for printing."""
         value = (((phi + np.pi) % (2*np.pi)) - np.pi) / np.pi
         if value == -1:
@@ -280,12 +287,15 @@ def phase_str(phi):
         # Terminate the last open regions
         if target_regions:
             target_regions[-1].append(t[-1])
-        for start, targets, end in target_regions:
+        for start, targets, end in target_regions:  # type: ignore
             q = targets[0]  # All targets have the same ref, so we pick
             ref = seq._phase_ref[basis][q]
             if end != seq._total_duration - 1 or 'measurement' not in data[ch]:
+                end = cast(int, end)
                 end += 1 / time_scale
-            for t_, delta in ref.changes(start, end, time_scale=time_scale):
+            for t_, delta in ref.changes(cast(int, start),
+                                         cast(int, end),
+                                         time_scale=time_scale):
                 conf = dict(linestyle='--', linewidth=1.5, color='black')
                 a.axvline(t_, **conf)
                 b.axvline(t_, **conf)

pulser/sequence.py

@@ -163,6 +163,7 @@ def __init__(self, register: Register, device: Device):
         # Checks if register is compatible with the device
         device.validate_register(register)
 
+        self._total_duration: int = 0
         self._register: Register = register
         self._device: Device = device
         self._in_xy: bool = False