Improvements to APT

instruments/abstract_instruments/comm/loopback_communicator.py

@@ -108,17 +108,21 @@ def read_raw(self, size=-1):
         :rtype: `bytes`
         """
         if self._stdin is not None:
-            if size >= 0:
+            if size == -1 or size is None:
+                result = bytes()
+                if self._terminator:
+                    while result.endswith(self._terminator.encode("utf-8")) is False:
+                        c = self._stdin.read(1)
+                        if c == b'':
+                            break
+                        result += c
+                    return result[:-len(self._terminator)]
+                return self._stdin.read(-1)
+
+            elif size >= 0:
                 input_var = self._stdin.read(size)
                 return bytes(input_var)
-            elif size == -1:
-                result = bytes()
-                while result.endswith(self._terminator.encode("utf-8")) is False:
-                    c = self._stdin.read(1)
-                    if c == b'':
-                        break
-                    result += c
-                return result[:-len(self._terminator)]
+
             else:
                 raise ValueError("Must read a positive value of characters.")
         else:

instruments/abstract_instruments/comm/serial_communicator.py

@@ -118,13 +118,23 @@ def read_raw(self, size=-1):
             return resp
         elif size == -1:
             result = bytes()
-            while result.endswith(self._terminator.encode("utf-8")) is False:
+            # If the terminator is empty, we can't use endswith, but must
+            # read as many bytes as are available.
+            # On the other hand, if terminator is nonempty, we can check
+            # that the tail end of the buffer matches it.
+            c = None
+            term = self._terminator.encode('utf-8') if self._terminator else None
+            while not (
+                    result.endswith(term)
+                    if term is not None else
+                    c == b''
+            ):
                 c = self._conn.read(1)
-                if c == b'':
+                if c == b'' and term is not None:
                     raise IOError("Serial connection timed out before reading "
                                   "a termination character.")
                 result += c
-            return result[:-len(self._terminator)]
+            return result[:-len(term)] if term is not None else result
         else:
             raise ValueError("Must read a positive value of characters.")
 

instruments/tests/test_comm/test_loopback.py

@@ -99,6 +99,25 @@ def test_loopbackcomm_read_raw_2char_terminator():
     assert mock_stdin.read.call_count == 5
 
 
+def test_loopbackcomm_read_raw_terminator_is_empty_string():
+    mock_stdin = mock.MagicMock()
+    mock_stdin.read.side_effect = [b"abc"]
+    comm = LoopbackCommunicator(stdin=mock_stdin)
+    comm._terminator = ""
+
+    assert comm.read_raw() == b"abc"
+    mock_stdin.read.assert_has_calls([mock.call(-1)])
+    assert mock_stdin.read.call_count == 1
+
+
+def test_loopbackcomm_read_raw_size_invalid():
+    with pytest.raises(ValueError):
+        mock_stdin = mock.MagicMock()
+        mock_stdin.read.side_effect = [b"abc"]
+        comm = LoopbackCommunicator(stdin=mock_stdin)
+        comm.read_raw(size=-2)
+
+
 def test_loopbackcomm_write_raw():
     mock_stdout = mock.MagicMock()
     comm = LoopbackCommunicator(stdout=mock_stdout)

instruments/tests/test_thorlabs/test_thorlabs_apt.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Module containing tests for the Thorlabs TC200
+"""
+
+# IMPORTS ####################################################################
+
+# pylint: disable=unused-import
+
+from __future__ import absolute_import
+
+import struct
+
+import pytest
+import quantities as pq
+
+import instruments as ik
+from instruments.thorlabs._packets import ThorLabsPacket, hw_info_data
+from instruments.thorlabs._cmds import ThorLabsCommands
+from instruments.tests import expected_protocol
+
+# TESTS ######################################################################
+
+# pylint: disable=protected-access,unused-argument
+
+
+def test_apt_hw_info():
+    with expected_protocol(
+        ik.thorlabs.ThorLabsAPT,
+        [
+            ThorLabsPacket(
+                message_id=ThorLabsCommands.HW_REQ_INFO,
+                param1=0x00, param2=0x00,
+                dest=0x50,
+                source=0x01,
+                data=None
+            ).pack()
+        ],
+        [
+            ThorLabsPacket(
+                message_id=ThorLabsCommands.HW_GET_INFO,
+                dest=0x01,
+                source=0x50,
+                data=hw_info_data.pack(
+                    # Serial number
+                    b'\x01\x02\x03\x04',
+                    # Model number
+                    "ABC-123".encode('ascii'),
+                    # HW type
+                    3,
+                    # FW version,
+                    0xa1, 0xa2, 0xa3,
+                    # Notes
+                    "abcdefg".encode('ascii'),
+                    # HW version
+                    42,
+                    # Mod state
+                    43,
+                    # Number of channels
+                    2
+                )
+            ).pack()
+        ],
+        sep=""
+    ) as apt:
+        # Check internal representations.
+        # NB: we shouldn't do this in some sense, but these fields
+        #     act as an API to the APT subclasses.
+        assert apt._hw_type == "Unknown type: 3"
+        assert apt._fw_version == "a1.a2.a3"
+        assert apt._notes == "abcdefg"
+        assert apt._hw_version == 42
+        assert apt._mod_state == 43
+
+        # Check external API.
+        assert apt.serial_number == '01020304'
+        assert apt.model_number == 'ABC-123'
+        assert apt.name == (
+            "ThorLabs APT Instrument model ABC-123, "
+            "serial 01020304 (HW version 42, FW version a1.a2.a3)"
+        )

instruments/thorlabs/_abstract.py

@@ -9,8 +9,13 @@
 from __future__ import absolute_import
 from __future__ import division
 
+import time
+
 from instruments.thorlabs import _packets
 from instruments.abstract_instruments.instrument import Instrument
+from instruments.util_fns import assume_units
+
+from quantities import second
 
 # CLASSES #####################################################################
 
@@ -35,10 +40,10 @@ def sendpacket(self, packet):
         :param packet: The thorlabs data packet that will be queried
         :type packet: `ThorLabsPacket`
         """
-        self.sendcmd(packet.pack())
+        self._file.write_raw(packet.pack())
 
     # pylint: disable=protected-access
-    def querypacket(self, packet, expect=None):
+    def querypacket(self, packet, expect=None, timeout=None, expect_data_len=None):
         """
         Sends a packet to the connected APT instrument, and waits for a packet
         in response. Optionally, checks whether the received packet type is
@@ -52,11 +57,40 @@ def querypacket(self, packet, expect=None):
             with the default value of `None` then no checking occurs.
         :type expect: `str` or `None`
 
+        :param timeout: Sets a timeout to wait before returning `None`, indicating
+            no packet was received. If the timeout is set to `None`, then the
+            timeout is inherited from the underlying communicator and no additional
+            timeout is added. If timeout is set to `False`, then this method waits
+            indefinitely. If timeout is set to a unitful quantity, then it is interpreted
+            as a time and used as the timeout value. Finally, if the timeout is a unitless
+            number (e.g. `float` or `int`), then seconds are assumed.
+
+        :param int expect_data_len: Number of bytes to expect as the
+            data for the returned packet.
+
         :return: Returns the response back from the instrument wrapped up in
-            a thorlabs packet
+            a ThorLabs APT packet, or None if no packet was received.
         :rtype: `ThorLabsPacket`
         """
-        resp = self.query(packet.pack())
+        t_start = time.time()
+
+        if timeout:
+            timeout = assume_units(timeout, second).rescale('second').magnitude
+
+        while True:
+            self._file.write_raw(packet.pack())
+            resp = self._file.read_raw(
+                expect_data_len + 6 # the header is six bytes.
+                if expect_data_len else
+                6
+            )
+            if resp or timeout is None:
+                break
+            else:
+                tic = time.time()
+                if tic - t_start > timeout:
+                    break
+
         if not resp:
             if expect is None:
                 return None
@@ -71,4 +105,5 @@ def querypacket(self, packet, expect=None):
             raise IOError("APT returned message ID {}, expected {}".format(
                 pkt._message_id, expect
             ))
+
         return pkt

instruments/thorlabs/_packets.py

@@ -15,6 +15,18 @@
 
 message_header_nopacket = struct.Struct('<HBBBB')
 message_header_wpacket = struct.Struct('<HHBB')
+hw_info_data = struct.Struct(
+    '<'    # Declare endianness.
+    '4s'   # serial_number
+    '8s'   # model_number
+    'H'    # hw_type_int
+    'BBBx' # fw_version
+    '48s'  # notes
+    '12x'  # padding
+    'H'    # hw_version
+    'H'    # mod_state
+    'H'    # n_channels
+)
 
 # CLASSES #####################################################################
 
@@ -166,7 +178,7 @@ def unpack(cls, bytes):
 
         # Check if 0x80 is set on header byte 4. If so, then this packet
         # has data.
-        if struct.unpack("B", header[4])[0] & 0x80:
+        if struct.unpack("B", header[4:5])[0] & 0x80:
             msg_id, length, dest, source = message_header_wpacket.unpack(
                 header)
             dest = dest ^ 0x80  # Turn off 0x80.