rm3100/ch347.py

"""I2C master over a CH347 USB adapter (Waveshare USB to UART/I2C/SPI/JTAG).

Knows nothing about any particular I2C device -- it only moves bytes to and
from a 7-bit address.

The CH347 in Mode 1 (UART1+I2C+SPI) presents three interfaces: two CDC-ACM
interfaces that the kernel binds for the UART, and interface 2, a vendor-class
interface carrying I2C/SPI/JTAG that no kernel driver claims. We talk to
interface 2 directly over libusb.

Framing follows the CH347 stream protocol as implemented by the
aystarik/ch347-i2c-spi-gpio Linux driver, whose USB id table entry
USB_DEVICE_INTERFACE_NUMBER(0x1a86, 0x55db, 0x02) matches this device.
"""

import time

import usb.core
import usb.util

VENDOR_ID = 0x1A86
PRODUCT_ID = 0x55DB
INTERFACE = 2

EP_OUT = 0x06
EP_IN = 0x86

TIMEOUT_MS = 1000

# Stream protocol opcodes.
CMD_STREAM = 0xAA  # start of an I2C command stream
CMD_END = 0x00  # end of stream
CMD_STA = 0x74  # emit START
CMD_STO = 0x75  # emit STOP
CMD_OUT = 0x80  # write; low 6 bits are the byte count
CMD_IN = 0xC0  # read; low 6 bits are the byte count
CMD_SET = 0x60  # set bus speed; low bits select the rate

# Bus speeds, as the low nibble of CMD_SET. The fastest is the default: bus
# speed does not set the sample rate -- the sensor does -- but it does set how
# long a read takes, and so how tightly an event on the far side of the bus can
# be timestamped. Measured host cost per sample is bus time plus ~0.6 ms of
# fixed USB round-trip latency: 2.13 ms at 100 kHz, 0.92 at 400, 0.79 at 750.
SPEED_20KHZ = 0
SPEED_100KHZ = 1
SPEED_400KHZ = 2
SPEED_750KHZ = 3

# The one place the rate constants are tied to their kHz figures. Callers that
# take a speed from a user or print one in a report go through this rather than
# keeping their own copy.
SPEEDS = {
    20: SPEED_20KHZ,
    100: SPEED_100KHZ,
    400: SPEED_400KHZ,
    750: SPEED_750KHZ,
}
DEFAULT_SPEED_KHZ = 750

# The count field is 6 bits, and a write also spends one byte on the address.
MAX_XFER = 0x3F

ACK = 1


class CH347I2C:
    """I2C master on a CH347 adapter.

    Every byte the adapter clocks out produces one status byte in the reply,
    where 1 means the slave ACKed. A read reply is one address-ACK byte
    followed by the payload.
    """

    def __init__(self, speed=SPEEDS[DEFAULT_SPEED_KHZ]):
        self._dev = usb.core.find(idVendor=VENDOR_ID, idProduct=PRODUCT_ID)
        if self._dev is None:
            raise IOError(
                f"No CH347 adapter found ({VENDOR_ID:04x}:{PRODUCT_ID:04x}). "
                "Is it plugged in and in Mode 1?"
            )

        # Deliberately no set_configuration(): the device is already configured,
        # and re-setting it would reset the CDC-ACM interfaces the kernel is
        # using for /dev/ttyACM*.
        try:
            usb.util.claim_interface(self._dev, INTERFACE)
        except usb.core.USBError as exc:
            raise IOError(
                f"Cannot claim CH347 interface {INTERFACE}: {exc}. "
                "If this is a permission error, run ./setup.sh."
            ) from exc

        self._claimed = True
        self.set_speed(speed)

    def close(self):
        if getattr(self, "_claimed", False):
            usb.util.release_interface(self._dev, INTERFACE)
            usb.util.dispose_resources(self._dev)
            self._claimed = False

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()
        return False

    def _xfer(self, out_bytes, in_len):
        """Send one command stream, then read in_len status/data bytes.

        Returns (data, mono, wall), the clocks read the instant the reply
        lands -- before any parsing or unwinding of the call stack. A caller
        timing an event on the far side of the bus (DRDY going high, say) wants
        that moment, not one several Python frames later: the frames add both a
        systematic lag and jitter from whatever the interpreter does between.

        The stamp travels with its own reply rather than being left on the
        instance, so it cannot go stale or be picked up by the wrong call.
        """
        written = self._dev.write(EP_OUT, bytes(out_bytes), TIMEOUT_MS)
        if written != len(out_bytes):
            raise IOError(
                f"Short USB write to CH347: sent {written} of {len(out_bytes)} bytes"
            )
        if in_len == 0:
            return b"", time.monotonic(), time.time()

        reply = bytes(self._dev.read(EP_IN, in_len, TIMEOUT_MS))
        mono, wall = time.monotonic(), time.time()
        if len(reply) != in_len:
            raise IOError(
                f"Short USB read from CH347: got {len(reply)} of {in_len} bytes"
            )
        return reply, mono, wall

    def set_speed(self, speed):
        """Select the I2C clock rate (one of the SPEED_* constants)."""
        if speed not in SPEEDS.values():
            raise ValueError(f"Invalid I2C speed {speed}, expected 0-3")
        self._xfer([CMD_STREAM, CMD_SET | speed, CMD_END], 0)

    def write(self, addr, data):
        """Write data (bytes) to a 7-bit address. Raises IOError on NACK."""
        data = bytes(data)
        # One byte of the transfer budget goes to the address.
        if len(data) > MAX_XFER - 1:
            raise ValueError(
                f"Write of {len(data)} bytes exceeds the CH347 limit of {MAX_XFER - 1}"
            )

        packet = [CMD_STREAM, CMD_STA, CMD_OUT | (len(data) + 1), addr << 1]
        packet += data
        packet += [CMD_STO, CMD_END]

        # One status byte per clocked-out byte: the address plus the payload.
        reply, _, _ = self._xfer(packet, len(data) + 1)
        if reply[0] != ACK:
            raise IOError(f"No ACK from I2C address 0x{addr:02x} on write")
        if any(byte != ACK for byte in reply[1:]):
            raise IOError(
                f"I2C address 0x{addr:02x} NACKed a data byte "
                f"(status {reply.hex(' ')})"
            )

    def read(self, addr, count):
        """Read count bytes from an address. Returns (data, mono, wall)."""
        if not 1 <= count <= MAX_XFER:
            raise ValueError(f"Read of {count} bytes outside 1..{MAX_XFER}")

        packet = [CMD_STREAM, CMD_STA, CMD_OUT | 1, (addr << 1) | 1]
        # All but the last byte are ACKed by us; the final bare CMD_IN NACKs to
        # tell the slave to stop.
        if count > 1:
            packet.append(CMD_IN | (count - 1))
        packet += [CMD_IN, CMD_STO, CMD_END]

        reply, mono, wall = self._xfer(packet, count + 1)
        if reply[0] != ACK:
            raise IOError(f"No ACK from I2C address 0x{addr:02x} on read")
        return reply[1:], mono, wall

    def write_read(self, addr, data, count):
        """Write data then read count bytes in one transaction.

        Returns (data, mono, wall) -- see _xfer() on why the stamp is returned
        rather than stored.

        Uses a repeated START rather than STOP-then-START, so the whole
        exchange is a single USB round trip instead of two. Measured at
        400 kHz this halves a one-byte register read (0.65 -> 0.35 ms), which
        matters because the host round trip, not the bus, sets the sample-rate
        ceiling.

        The RM3100's manual draws its register reads with a STOP between the
        pointer write and the read (sections 4.5.2, 5.8.4), but the part
        accepts a repeated START -- verified against REVID and the measurement
        registers.
        """
        data = bytes(data)
        if not data:
            raise ValueError("write_read needs at least one byte to write")
        if len(data) > MAX_XFER - 1:
            raise ValueError(f"Write of {len(data)} bytes exceeds the CH347 limit")
        if not 1 <= count <= MAX_XFER:
            raise ValueError(f"Read of {count} bytes outside 1..{MAX_XFER}")

        packet = [CMD_STREAM, CMD_STA, CMD_OUT | (len(data) + 1), addr << 1]
        packet += data
        packet += [CMD_STA, CMD_OUT | 1, (addr << 1) | 1]
        if count > 1:
            packet.append(CMD_IN | (count - 1))
        packet += [CMD_IN, CMD_STO, CMD_END]

        # One status byte per clocked-out byte: the write address, the payload,
        # and the read address -- then the payload itself.
        acks = len(data) + 2
        reply, mono, wall = self._xfer(packet, acks + count)
        if any(byte != ACK for byte in reply[:acks]):
            raise IOError(
                f"I2C address 0x{addr:02x} NACKed during combined transfer "
                f"(status {reply[:acks].hex(' ')})")
        return reply[acks:], mono, wall

    def probe(self, addr):
        """Return True if a device ACKs its address. A NACK is not an error."""
        packet = [CMD_STREAM, CMD_STA, CMD_OUT | 1, addr << 1, CMD_STO, CMD_END]
        try:
            return self._xfer(packet, 1)[0][0] == ACK
        except IOError:
            return False

    def scan(self, first=0x08, last=0x77):
        """Return the addresses in [first, last] that respond."""
        return [addr for addr in range(first, last + 1) if self.probe(addr)]