python: multi_usrp: Fix issues in send_waveform()

- Like with RX, this now allows passing in stream time and existing
  streamer
- There was no EOB being sent at the end (now there is)
- Fixed some linter issues

2 files changed, 36 insertions, 19 deletions

diff --git a/host/examples/python/tx_waveforms.py b/host/examples/python/tx_waveforms.py
index 58b6472f7..2e6f5ecde 100755
--- a/host/examples/python/tx_waveforms.py
+++ b/host/examples/python/tx_waveforms.py
@@ -12,9 +12,9 @@ import argparse
 import numpy as np
 import uhd
 
-waveforms = {
+WAVEFORMS = {
     "sine": lambda n, tone_offset, rate: np.exp(n * 2j * np.pi * tone_offset / rate),
-    "square": lambda n, tone_offset, rate: np.sign(waveforms["sine"](n, tone_offset, rate)),
+    "square": lambda n, tone_offset, rate: np.sign(WAVEFORMS["sine"](n, tone_offset, rate)),
     "const": lambda n, tone_offset, rate: 1 + 1j,
     "ramp": lambda n, tone_offset, rate:
             2*(n*(tone_offset/rate) - np.floor(float(0.5 + n*(tone_offset/rate))))
@@ -26,7 +26,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("-a", "--args", default="", type=str)
     parser.add_argument(
-        "-w", "--waveform", default="sine", choices=waveforms.keys(), type=str)
+        "-w", "--waveform", default="sine", choices=WAVEFORMS.keys(), type=str)
     parser.add_argument("-f", "--freq", type=float, required=True)
     parser.add_argument("-r", "--rate", default=1e6, type=float)
     parser.add_argument("-d", "--duration", default=5.0, type=float)
@@ -44,10 +44,10 @@ def main():
     if not isinstance(args.channels, list):
         args.channels = [args.channels]
     data = np.array(
-        list(map(lambda n: args.wave_ampl * waveforms[args.waveform](n, args.wave_freq, args.rate),
-            np.arange(
-                int(10 * np.floor(args.rate / args.wave_freq)),
-                dtype=np.complex64))),
+        list(map(lambda n: args.wave_ampl * WAVEFORMS[args.waveform](n, args.wave_freq, args.rate),
+                 np.arange(
+                     int(10 * np.floor(args.rate / args.wave_freq)),
+                     dtype=np.complex64))),
         dtype=np.complex64)  # One period
 
     usrp.send_waveform(data, args.duration, args.freq, args.rate,
diff --git a/host/python/uhd/usrp/multi_usrp.py b/host/python/uhd/usrp/multi_usrp.py
index a025b10d1..f9aa1a857 100644
--- a/host/python/uhd/usrp/multi_usrp.py
+++ b/host/python/uhd/usrp/multi_usrp.py
@@ -87,7 +87,7 @@ class MultiUSRP(lib.usrp.multi_usrp):
             streamer.issue_stream_cmd(stream_cmd)
         def _stop_stream(streamer):
             """
-            Issue the start-stream command and flush the queue.
+            Issue the stop-stream command and flush the queue.
             """
             metadata = lib.types.rx_metadata()
             stream_cmd = lib.types.stream_cmd(lib.types.stream_mode.stop_cont)
@@ -132,7 +132,9 @@ class MultiUSRP(lib.usrp.multi_usrp):
                       freq,
                       rate=1e6,
                       channels=(0,),
-                      gain=10):
+                      gain=10,
+                      start_time=None,
+                      streamer=None):
         """
         TX a finite number of samples from the USRP
         :param waveform_proto: numpy array of samples to TX
@@ -141,42 +143,57 @@ class MultiUSRP(lib.usrp.multi_usrp):
         :param rate: TX sample rate (Hz)
         :param channels: list of channels to TX on
         :param gain: TX gain (dB)
+        :param start_time: A valid TimeSpec object with the starting time. If
+                           None, then streaming starts immediately.
+        :param streamer: A TX streamer object. If None, this function will create
+                         one locally and attempt to destroy it afterwards.
         :return: the number of transmitted samples
         """
+        def _config_streamer(streamer):
+            """
+            Set up the correct streamer
+            """
+            if streamer is None:
+                st_args = lib.usrp.stream_args("fc32", "sc16")
+                st_args.channels = channels
+                streamer = super(MultiUSRP, self).get_tx_stream(st_args)
+            return streamer
+        ## And go!
         for chan in channels:
             super(MultiUSRP, self).set_tx_rate(rate, chan)
             super(MultiUSRP, self).set_tx_freq(lib.types.tune_request(freq), chan)
             super(MultiUSRP, self).set_tx_gain(gain, chan)
 
-        st_args = lib.usrp.stream_args("fc32", "sc16")
-        st_args.channels = channels
-
-        streamer = super(MultiUSRP, self).get_tx_stream(st_args)
+        # Configure streamer
+        streamer = _config_streamer(streamer)
+        # Set up buffers and counters
         buffer_samps = streamer.get_max_num_samps()
         proto_len = waveform_proto.shape[-1]
-
         if proto_len < buffer_samps:
             waveform_proto = np.tile(waveform_proto,
                                      (1, int(np.ceil(float(buffer_samps)/proto_len))))
             proto_len = waveform_proto.shape[-1]
-
-        metadata = lib.types.tx_metadata()
         send_samps = 0
         max_samps = int(np.floor(duration * rate))
-
         if len(waveform_proto.shape) == 1:
             waveform_proto = waveform_proto.reshape(1, waveform_proto.size)
         if waveform_proto.shape[0] < len(channels):
             waveform_proto = np.tile(waveform_proto[0], (len(channels), 1))
-
+        # Now stream
+        metadata = lib.types.tx_metadata()
+        if start_time is not None:
+            metadata.time_spec = start_time
         while send_samps < max_samps:
             real_samps = min(proto_len, max_samps-send_samps)
             if real_samps < proto_len:
+                print(waveform_proto[:, :real_samps])
                 samples = streamer.send(waveform_proto[:, :real_samps], metadata)
             else:
                 samples = streamer.send(waveform_proto, metadata)
             send_samps += samples
-
+        # Send EOB to terminate Tx
+        metadata.end_of_burst = True
+        streamer.send(np.zeros((len(channels), 1), dtype=np.complex64), metadata)
         # Help the garbage collection
         streamer = None
         return send_samps