Source code for pandablocks.connections

import struct
import sys
import xml.etree.ElementTree as ET
from collections import deque
from dataclasses import dataclass
from typing import Any, Callable, Deque, Iterator, List, Optional, Tuple

import numpy as np

from ._exchange import Exchange, ExchangeGenerator, Exchanges
from .commands import Command, CommandException
from .responses import (
    Data,
    EndData,
    EndReason,
    FieldCapture,
    FrameData,
    ReadyData,
    StartData,
)

# Define the public API of this module
__all__ = [
    "NeedMoreData",
    "NoContextAvailable",
    "Buffer",
    "ControlConnection",
    "DataConnection",
]

# The name of the samples field used for averaging unscaled fields
SAMPLES_FIELD = "PCAP.SAMPLES.Value"



[docs]
class NeedMoreData(Exception):
    """Raised if the `Buffer` isn't full enough to return the requested bytes"""




[docs]
class NoContextAvailable(Exception):
    """Raised if there were no contexts available for this connection.
    This may result from calling `ControlConnection.receive_bytes()` without calling
    `ControlConnection.send()`, or if there were unmatched sends/receives"""




[docs]
class Buffer:
    """Byte storage that provides line reader and bytes reader interfaces.
    For example::

        buf = Buffer()
        buf += bytes_from_server
        line = buf.read_line()  # raises NeedMoreData if no line
        for line in buf:
            pass
        bytes = buf.read_bytes(50)  # raises NeedMoreData if not enough bytes
    """

    def __init__(self):
        self._buf = bytearray()

    def __iadd__(self, byteslike: bytes):
        """Add some data from the server"""
        self._buf += byteslike
        return self

    def _extract_frame(self, num_to_extract, num_to_discard=0) -> bytearray:
        # extract num_to_extract bytes from the start of the buffer
        frame = self._buf[:num_to_extract]
        # Update the buffer in place, to take advantage of bytearray's
        # optimized delete-from-beginning feature.
        del self._buf[: num_to_extract + num_to_discard]
        return frame


[docs]
    def read_bytes(self, num: int) -> bytearray:
        """Read and pop num bytes from the beginning of the buffer, raising
        `NeedMoreData` if the buffer isn't full enough to do so"""
        if num > len(self._buf):
            raise NeedMoreData()
        else:
            return self._extract_frame(num)



[docs]
    def peek_bytes(self, num: int) -> bytearray:
        """Read but do not pop num bytes from the beginning of the buffer,
        raising `NeedMoreData` if the buffer isn't full enough to do so"""
        if num > len(self._buf):
            raise NeedMoreData()
        else:
            return self._buf[:num]



[docs]
    def read_line(self):
        """Read and pop a newline terminated line (without terminator)
        from the beginning of the buffer, raising `NeedMoreData` if the
        buffer isn't full enough to do so"""
        idx = self._buf.find(b"\n")
        if idx < 0:
            raise NeedMoreData()
        else:
            return self._extract_frame(idx, num_to_discard=1)


    def __iter__(self):
        return self

    def __next__(self) -> bytes:
        try:
            return self.read_line()
        except NeedMoreData:
            raise StopIteration()



@dataclass
class _ExchangeContext:
    #: The exchange we should be filling
    exchange: Exchange
    #: The command that produced it
    command: Command
    #: If this was the last in the list, the generator to call next
    generator: Optional[ExchangeGenerator[Any]] = None

    def exception(self, e: Exception) -> CommandException:
        """Return a `CommandException` with the sent and received strings
        in the text"""
        msg = f"{self.command} ->"
        if self.exchange.is_multiline:
            for line in self.exchange.multiline:
                msg += "\n    " + line
        else:
            msg += " " + self.exchange.line
        if e.args:
            msg += f"\n{type(e).__name__}:{e}"
        return CommandException(msg).with_traceback(e.__traceback__)



[docs]
class ControlConnection:
    """Sans-IO connection to control port of PandA TCP server, supporting a
    Command based interface. For example::

        cc = ControlConnection()
        # Connection says what bytes should be sent to execute command
        to_send = cc.send(command)
        socket.sendall(to_send)
        while True:
            # Repeatedly process bytes from the PandA
            received = socket.recv()
            # Sending any subsequent bytes to be sent back to the PandA
            to_send = cc.receive_bytes(received)
            socket.sendall(to_send)
            # And processing the produced responses
            for command, response in cc.responses()
                do_something_with(response)
    """

    def __init__(self) -> None:
        self._buf = Buffer()
        self._lines: List[str] = []
        self._contexts: Deque[_ExchangeContext] = deque()
        self._responses: Deque[Tuple[Command, Any]] = deque()

    def _update_contexts(self, lines: List[str], is_multiline=False) -> bytes:
        to_send = b""
        if len(self._contexts) == 0:
            raise NoContextAvailable()
        context = self._contexts.popleft()
        # Update the exchange with what we've got
        context.exchange.received = lines
        context.exchange.is_multiline = is_multiline
        # If we're given a generator to run then do so
        if context.generator:
            try:
                # Return the bytes from sending the next bit of the command
                exchanges = next(context.generator)
            except StopIteration as e:
                # Command complete, store the result
                self._responses.append((context.command, e.value))
            except Exception as e:
                # Command failed, store an exception
                self._responses.append((context.command, context.exception(e)))
            else:
                to_send = b"".join(
                    self._bytes_from_exchanges(
                        exchanges, context.command, context.generator
                    )
                )
        return to_send

    def _bytes_from_exchanges(
        self, exchanges: Exchanges, command: Command, generator: ExchangeGenerator[Any]
    ) -> Iterator[bytes]:
        if not isinstance(exchanges, list):
            exchanges = [exchanges]
        # No Exchanges when a Command's yield is empty e.g. unexpected/unparseable data
        # received from PandA
        if len(exchanges) == 0:
            return
        for ex in exchanges:
            context = _ExchangeContext(ex, command)
            self._contexts.append(context)
            text = "\n".join(ex.to_send) + "\n"
            yield text.encode()
        # The last exchange gets the generator so it triggers the next thing to send
        context.generator = generator


[docs]
    def receive_bytes(self, received: bytes) -> bytes:
        """Tell the connection that you have received some bytes off the network.
        Parse these into high level responses which are yielded back by `responses`.
        Return any bytes to send back"""
        self._buf += received
        is_multiline = bool(self._lines)
        to_send = b""
        for line_b in self._buf:
            line = line_b.decode()
            if not is_multiline:
                # Check if we need to switch to multiline mode
                is_multiline = line.startswith("!") or line == "."
            if is_multiline:
                # Add a new line to the buffer
                self._lines.append(line)
                if line == ".":
                    # End of multiline mode, return what we've got
                    to_send += self._update_contexts(self._lines, is_multiline)
                    self._lines = []
                    is_multiline = False
                else:
                    # Check a correctly formatted response
                    assert line.startswith("!"), (
                        "Multiline response %r doesn't start with !" % line
                    )
            else:
                # Single line mode
                assert not self._lines, (
                    "Multiline response %s not terminated" % self._lines
                )
                to_send += self._update_contexts([line])
        return to_send



[docs]
    def responses(self) -> Iterator[Tuple[Command, Any]]:
        """Get the (command, response) tuples generated as part of the last
        receive_bytes"""
        while self._responses:
            yield self._responses.popleft()



[docs]
    def send(self, command: Command) -> bytes:
        """Tell the connection you want to send an event, and it will return
        some bytes to send down the network
        """
        # If not given a partially run generator, start one here
        generator = command.execute()
        exchanges = next(generator)
        to_send = b"".join(self._bytes_from_exchanges(exchanges, command, generator))
        return to_send





[docs]
class DataConnection:
    """Sans-IO connection to data port of PandA TCP server, supporting an
    flushable iterator interface. For example::

        dc = DataConnection()
        # Single connection string to send
        to_send = dc.connect()
        socket.sendall(to_send)
        while True:
            # Repeatedly process bytes from the PandA looking for data
            received = socket.recv()
            for data in dc.receive_bytes(received):
                do_something_with(data)
    """

    def __init__(self) -> None:
        # TODO: could support big endian, but are there any systems out there?
        assert sys.byteorder == "little", "PandA sends data little endian"
        # Store of bytes received so far to parse in the handlers
        self._buf = Buffer()
        # Header text from PandA with field info
        self._header = ""
        # The next parsing handler that should be called if there is data in buffer
        self._next_handler: Optional[Callable[[], Optional[Iterator[Data]]]] = None
        # numpy dtype of produced FrameData
        self._frame_dtype = None
        # frame data that has been received but not flushed yet
        self._partial_data = bytearray()
        # whether to flush after every frame
        self._flush_every_frame = False

    def _handle_connected(self):
        # Get the response from connect()
        line = self._buf.read_line()
        assert line == b"OK", f"Expected OK, got {line!r}"
        yield ReadyData()
        self._next_handler = self._handle_header_start

    def _handle_header_start(self):
        # Discard lines until we see header start tag
        line = self._buf.read_line()
        if line == b"<header>":
            self._header = line
            self._next_handler = self._handle_header_body

    def _handle_header_body(self):
        # Accumumlate header until the end tag, then parese and return
        line = self._buf.read_line()
        self._header += line
        if line == b"</header>":
            fields = []
            root = ET.fromstring(self._header)
            for field in root.find("fields"):
                fields.append(
                    FieldCapture(
                        name=str(field.get("name")),
                        type=np.dtype(field.get("type")),
                        capture=str(field.get("capture")),
                        scale=float(field.get("scale", 1)),
                        offset=float(field.get("offset", 0)),
                        units=str(field.get("units", "")),
                    )
                )
            data = root.find("data")
            sample_bytes = int(data.get("sample_bytes"))
            if sample_bytes - sum(f.type.itemsize for f in fields) == 4:
                # In raw mode with panda-server < 2.1 samples wasn't
                # put in if not specifically requested, but was still
                # sent
                name, capture = SAMPLES_FIELD.rsplit(".", maxsplit=1)
                fields.insert(
                    0,
                    FieldCapture(name, np.dtype("uint32"), capture),
                )
            self._frame_dtype = np.dtype(
                [(f"{f.name}.{f.capture}", f.type) for f in fields]
            )
            yield StartData(
                fields=fields,
                missed=int(data.get("missed")),
                process=str(data.get("process")),
                format=str(data.get("format")),
                sample_bytes=sample_bytes,
            )
            self._next_handler = self._handle_header_end

    def _handle_header_end(self):
        # Discard the newline at the end of the header
        assert self._buf.read_bytes(1) == b"\n", "Expected newline at end of header"
        self._next_handler = self._handle_data_start

    def _handle_data_start(self):
        # Handle "BIN " or "END "
        bytes = self._buf.read_bytes(4)
        if bytes == b"BIN ":
            self._next_handler = self._handle_data_frame
        elif bytes == b"END ":
            self._next_handler = self._handle_data_end
        else:
            raise ValueError(f"Bad data '{bytes}'")

    def _handle_data_frame(self):
        # Handle a whole data frame
        # Peek message length as uint32 LE
        # length = len("BIN " + 4_bytes_encoded_length + data)
        length = struct.unpack("<I", self._buf.peek_bytes(4))[0]
        # we already read "BIN ", so read the rest
        data = self._buf.read_bytes(length - 4)[4:]
        self._partial_data += data
        # if told to flush now, then yield what we have
        if self._flush_every_frame:
            yield from self.flush()
        self._next_handler = self._handle_data_start

    def _handle_data_end(self):
        # Handle the end reason
        samples, reason = self._buf.read_line().split(maxsplit=1)
        reason_enum = EndReason(reason.decode())
        # Flush whatever is not already flushed
        yield from self.flush()
        yield EndData(samples=int(samples), reason=reason_enum)
        self._next_handler = self._handle_header_start


[docs]
    def receive_bytes(self, received: bytes, flush_every_frame=True) -> Iterator[Data]:
        """Tell the connection that you have received some bytes off the network.
        Parse these into Data structures and yield them back.

        Args:
            received: the bytes you received from the socket
            flush_every_frame: Whether to flush `FrameData` as soon as received.
                If False then they will only be sent if `flush` is called or
                end of acquisition reached
        """
        assert self._next_handler, "Connect not called"
        self._flush_every_frame = flush_every_frame
        self._buf += received
        while True:
            # Each of these handlers should call read at most once, so that
            # if we don't have enough data we don't lose partial data
            try:
                ret = self._next_handler()
                if ret:
                    # This is an iterator of Data objects
                    yield from ret
            except NeedMoreData:
                break



[docs]
    def flush(self) -> Iterator[FrameData]:
        """If there is a partial data frame, pop and yield it"""
        if self._partial_data:
            # Make a numpy array wrapper to the bytearray, no copying here
            data = np.frombuffer(self._partial_data, self._frame_dtype)
            # Make a new bytearray, numpy view will keep the reference
            # to the old one so can't clear it in place
            self._partial_data = bytearray()
            yield FrameData(data)



[docs]
    def connect(self, scaled: bool) -> bytes:
        """Return the bytes that need to be sent on connection"""
        assert not self._next_handler, "Can't connect while midway through collection"
        self._next_handler = self._handle_connected
        if scaled:
            return b"XML FRAMED SCALED\n"
        else:
            return b"XML FRAMED RAW\n"