From 565537ec737c845d4bedec83b8a2b83fd0aa61a9 Mon Sep 17 00:00:00 2001
From: Santosh Bala <santmartian@gmail.com>
Date: Tue, 28 Jul 2026 16:53:29 +0530
Subject: [PATCH] feat: add Pico SCPI transport

---
 pslab/__init__.py                 |   2 +
 pslab/pico/__init__.py            |  21 ++
 pslab/pico/device.py              |  93 +++++++++
 pslab/pico/transport.py           | 336 ++++++++++++++++++++++++++++++
 tests/test_pico_scpi_transport.py | 126 +++++++++++
 5 files changed, 578 insertions(+)
 create mode 100644 pslab/pico/__init__.py
 create mode 100644 pslab/pico/device.py
 create mode 100644 pslab/pico/transport.py
 create mode 100644 tests/test_pico_scpi_transport.py

diff --git a/pslab/__init__.py b/pslab/__init__.py
index fde015c..e809319 100644
--- a/pslab/__init__.py
+++ b/pslab/__init__.py
@@ -5,12 +5,14 @@
 from pslab.instrument.oscilloscope import Oscilloscope
 from pslab.instrument.power_supply import PowerSupply
 from pslab.instrument.waveform_generator import PWMGenerator, WaveformGenerator
+from pslab.pico import PicoDevice
 from pslab.sciencelab import ScienceLab
 
 __all__ = (
     "LogicAnalyzer",
     "Multimeter",
     "Oscilloscope",
+    "PicoDevice",
     "PowerSupply",
     "PWMGenerator",
     "WaveformGenerator",
diff --git a/pslab/pico/__init__.py b/pslab/pico/__init__.py
new file mode 100644
index 0000000..e8026b7
--- /dev/null
+++ b/pslab/pico/__init__.py
@@ -0,0 +1,21 @@
+"""PSLab Pico SCPI support."""
+
+from pslab.pico.device import PicoDevice
+from pslab.pico.transport import (
+    PicoTransport,
+    PicoUsbTransport,
+    PicoWifiTransport,
+    ScpiClient,
+    ScpiError,
+    ScpiTimeoutError,
+)
+
+__all__ = (
+    "PicoDevice",
+    "PicoTransport",
+    "PicoUsbTransport",
+    "PicoWifiTransport",
+    "ScpiClient",
+    "ScpiError",
+    "ScpiTimeoutError",
+)
diff --git a/pslab/pico/device.py b/pslab/pico/device.py
new file mode 100644
index 0000000..0cec6b7
--- /dev/null
+++ b/pslab/pico/device.py
@@ -0,0 +1,93 @@
+"""High-level PSLab Pico device entry point."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pslab.pico.transport import PicoTransport, PicoUsbTransport, PicoWifiTransport, ScpiClient
+
+
+class PicoDevice:
+    """Connected PSLab Pico firmware device.
+
+    The Pico firmware speaks SCPI over USB CDC or over the ESP32-C3 TCP bridge.
+    This object owns the shared SCPI client used by future Pico instruments.
+    """
+
+    def __init__(self, client: ScpiClient) -> None:
+        self.scpi = client
+
+    @classmethod
+    def usb(
+        cls,
+        port: Optional[str] = None,
+        baudrate: int = 115200,
+        timeout: Optional[float] = 2.0,
+    ) -> "PicoDevice":
+        """Create a Pico device using USB CDC."""
+
+        return cls(ScpiClient(PicoUsbTransport(port, baudrate, timeout)))
+
+    @classmethod
+    def wifi(
+        cls,
+        host: str,
+        port: int = PicoWifiTransport.DEFAULT_PORT,
+        timeout: Optional[float] = 2.0,
+    ) -> "PicoDevice":
+        """Create a Pico device using the ESP TCP SCPI bridge."""
+
+        return cls(ScpiClient(PicoWifiTransport(host, port, timeout)))
+
+    @classmethod
+    def from_transport(cls, transport: PicoTransport) -> "PicoDevice":
+        """Create a Pico device from a custom transport."""
+
+        return cls(ScpiClient(transport))
+
+    def connect(self) -> None:
+        """Open the underlying transport."""
+
+        self.scpi.connect()
+
+    def close(self) -> None:
+        """Close the underlying transport."""
+
+        self.scpi.close()
+
+    def identify(self) -> str:
+        """Return the firmware identification string."""
+
+        return self.scpi.identify()
+
+    def reset(self) -> None:
+        """Reset all instrument state in the firmware."""
+
+        self.scpi.reset()
+
+    def clear_status(self) -> None:
+        """Clear the SCPI status and error queues."""
+
+        self.scpi.clear_status()
+
+    def set_transport(self, mode: str) -> None:
+        """Select capture data transport: ``USB``, ``WIFI``, or ``AUTO``."""
+
+        self.scpi.set_transport(mode)
+
+    def get_transport(self) -> str:
+        """Return the selected capture data transport."""
+
+        return self.scpi.get_transport()
+
+    def wifi_status(self):
+        """Return Wi-Fi transport status counters."""
+
+        return self.scpi.wifi_status()
+
+    def __enter__(self) -> "PicoDevice":
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc, traceback) -> None:
+        self.close()
diff --git a/pslab/pico/transport.py b/pslab/pico/transport.py
new file mode 100644
index 0000000..1500539
--- /dev/null
+++ b/pslab/pico/transport.py
@@ -0,0 +1,336 @@
+"""SCPI transport helpers for PSLab Pico firmware."""
+
+from __future__ import annotations
+
+import socket
+from typing import List, Optional, Tuple
+
+import serial
+from serial.tools import list_ports
+
+
+class ScpiError(RuntimeError):
+    """Raised when the device returns a SCPI error response."""
+
+
+class ScpiTimeoutError(TimeoutError):
+    """Raised when a SCPI response is not received before timeout."""
+
+
+class PicoTransport:
+    """Small byte-oriented transport interface used by :class:`ScpiClient`."""
+
+    timeout: Optional[float]
+
+    def connect(self) -> None:
+        """Open the transport."""
+        raise NotImplementedError
+
+    def close(self) -> None:
+        """Close the transport."""
+        raise NotImplementedError
+
+    def read(self, size: int) -> bytes:
+        """Read up to ``size`` bytes."""
+        raise NotImplementedError
+
+    def readline(self) -> bytes:
+        """Read one newline-terminated line."""
+        raise NotImplementedError
+
+    def write(self, data: bytes) -> int:
+        """Write bytes to the device."""
+        raise NotImplementedError
+
+    def flush(self) -> None:
+        """Flush pending output bytes if the backend supports it."""
+
+    def reset_input_buffer(self) -> None:
+        """Discard unread input bytes if the backend supports it."""
+
+
+class PicoUsbTransport(PicoTransport):
+    """USB CDC transport for PSLab Pico SCPI firmware."""
+
+    USB_IDS = ((0xCAFE, 0x4010),)
+
+    def __init__(
+        self,
+        port: Optional[str] = None,
+        baudrate: int = 115200,
+        timeout: Optional[float] = 2.0,
+        serial_instance=None,
+    ) -> None:
+        self.port = port
+        self.baudrate = baudrate
+        self.timeout = timeout
+        self._serial = serial_instance
+
+    @classmethod
+    def find_ports(cls) -> List[str]:
+        """Return candidate serial ports for PSLab Pico devices."""
+
+        ports = []
+        for port in list_ports.comports():
+            if (port.vid, port.pid) in cls.USB_IDS:
+                ports.append(port.device)
+                continue
+            if port.product and port.product.startswith("PSLab Pico"):
+                ports.append(port.device)
+        return ports
+
+    @classmethod
+    def find_port(cls) -> str:
+        """Return the only detected PSLab Pico serial port."""
+
+        ports = cls.find_ports()
+        if not ports:
+            raise ConnectionError("PSLab Pico USB CDC port not found.")
+        if len(ports) > 1:
+            raise ConnectionError(f"Multiple PSLab Pico ports found: {ports}")
+        return ports[0]
+
+    @property
+    def is_open(self) -> bool:
+        return bool(self._serial is not None and getattr(self._serial, "is_open", True))
+
+    def connect(self) -> None:
+        if self.is_open:
+            return
+        if self.port is None:
+            self.port = self.find_port()
+        self._serial = serial.Serial(
+            self.port,
+            baudrate=self.baudrate,
+            timeout=self.timeout,
+            write_timeout=self.timeout,
+        )
+
+    def close(self) -> None:
+        if self._serial is not None:
+            self._serial.close()
+
+    def read(self, size: int) -> bytes:
+        self.connect()
+        return self._serial.read(size)
+
+    def readline(self) -> bytes:
+        self.connect()
+        return self._serial.readline()
+
+    def write(self, data: bytes) -> int:
+        self.connect()
+        return self._serial.write(data)
+
+    def flush(self) -> None:
+        if self._serial is not None:
+            self._serial.flush()
+
+    def reset_input_buffer(self) -> None:
+        if self._serial is not None and hasattr(self._serial, "reset_input_buffer"):
+            self._serial.reset_input_buffer()
+
+
+class PicoWifiTransport(PicoTransport):
+    """TCP transport for the ESP32-C3 PSLab Pico SCPI bridge."""
+
+    DEFAULT_PORT = 5006
+
+    def __init__(
+        self,
+        host: str,
+        port: int = DEFAULT_PORT,
+        timeout: Optional[float] = 2.0,
+    ) -> None:
+        self.host = host
+        self.port = int(port)
+        self.timeout = timeout
+        self._socket = None
+        self._rx_buffer = bytearray()
+
+    @property
+    def is_open(self) -> bool:
+        return self._socket is not None
+
+    def connect(self) -> None:
+        if self._socket is not None:
+            return
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(self.timeout)
+        sock.connect((self.host, self.port))
+        self._socket = sock
+
+    def close(self) -> None:
+        if self._socket is not None:
+            self._socket.close()
+            self._socket = None
+        self._rx_buffer.clear()
+
+    def read(self, size: int) -> bytes:
+        self.connect()
+        while len(self._rx_buffer) < size:
+            chunk = self._socket.recv(max(1, size - len(self._rx_buffer)))
+            if not chunk:
+                raise ScpiTimeoutError("SCPI TCP connection closed.")
+            self._rx_buffer.extend(chunk)
+        data = bytes(self._rx_buffer[:size])
+        del self._rx_buffer[:size]
+        return data
+
+    def readline(self) -> bytes:
+        self.connect()
+        while b"\n" not in self._rx_buffer:
+            chunk = self._socket.recv(4096)
+            if not chunk:
+                raise ScpiTimeoutError("SCPI TCP connection closed.")
+            self._rx_buffer.extend(chunk)
+        end = self._rx_buffer.index(b"\n") + 1
+        line = bytes(self._rx_buffer[:end])
+        del self._rx_buffer[:end]
+        return line
+
+    def write(self, data: bytes) -> int:
+        self.connect()
+        self._socket.sendall(data)
+        return len(data)
+
+    def reset_input_buffer(self) -> None:
+        self._rx_buffer.clear()
+
+
+class ScpiClient:
+    """Line-oriented SCPI client for PSLab Pico firmware."""
+
+    def __init__(self, transport: PicoTransport) -> None:
+        self.transport = transport
+
+    def connect(self) -> None:
+        self.transport.connect()
+
+    def close(self) -> None:
+        self.transport.close()
+
+    def command(self, command: str) -> None:
+        """Send a command that does not return a response."""
+
+        if command.rstrip().endswith("?"):
+            raise ValueError("Use query() or query_block() for SCPI queries.")
+        self._write_line(command)
+
+    def query(self, command: str) -> str:
+        """Send a query and return one decoded response line."""
+
+        self._write_line(command)
+        line = self.transport.readline()
+        if not line:
+            raise ScpiTimeoutError(f"Timed out waiting for response to {command!r}.")
+        return line.decode("ascii", errors="replace").strip()
+
+    def query_block(self, command: str) -> bytes:
+        """Send a query and return a SCPI definite-length block payload."""
+
+        self._write_line(command)
+        return self.read_block()
+
+    def read_block(self) -> bytes:
+        """Read one SCPI definite-length arbitrary block payload."""
+
+        marker = self._read_exact(1)
+        if marker != b"#":
+            rest = self.transport.readline()
+            message = (marker + rest).decode("ascii", errors="replace").strip()
+            raise ScpiError(message)
+
+        digit_count_text = self._read_exact(1)
+        try:
+            digit_count = int(digit_count_text.decode("ascii"))
+            byte_count = int(self._read_exact(digit_count).decode("ascii"))
+        except ValueError as exc:
+            raise ScpiError("Invalid SCPI arbitrary block header.") from exc
+
+        payload = self._read_exact(byte_count)
+        self._drain_block_terminator()
+        return payload
+
+    def identify(self) -> str:
+        """Return the firmware identification string."""
+
+        return self.query("*IDN?")
+
+    def reset(self) -> None:
+        """Reset instrument state."""
+
+        self.command("*RST")
+
+    def clear_status(self) -> None:
+        """Clear SCPI status/error state."""
+
+        self.command("*CLS")
+
+    def error_count(self) -> int:
+        """Return the number of queued SCPI errors."""
+
+        return int(self.query("SYST:ERR:COUN?").strip())
+
+    def system_error(self) -> Tuple[int, str]:
+        """Pop one SCPI error from the device error queue."""
+
+        response = self.query("SYST:ERR?")
+        code_text, _, message = response.partition(",")
+        return int(code_text), message.strip().strip('"')
+
+    def drain_errors(self) -> List[Tuple[int, str]]:
+        """Pop all queued SCPI errors."""
+
+        errors = []
+        while True:
+            error = self.system_error()
+            errors.append(error)
+            if error[0] == 0:
+                return errors
+
+    def set_transport(self, mode: str) -> None:
+        """Select firmware capture transport: ``USB``, ``WIFI``, or ``AUTO``."""
+
+        normalized = mode.strip().upper()
+        if normalized not in ("USB", "WIFI", "WIRELESS", "AUTO"):
+            raise ValueError("mode must be USB, WIFI, WIRELESS, or AUTO.")
+        self.command(f"COMM:TRAN {normalized}")
+
+    def get_transport(self) -> str:
+        """Return selected firmware capture transport."""
+
+        return self.query("COMM:TRAN?")
+
+    def wifi_status(self) -> Tuple[bool, int, int, int]:
+        """Return Wi-Fi transport counters.
+
+        Returns
+        -------
+        effective, sent_frames, dropped_frames, timeouts
+        """
+
+        fields = self.query("COMM:WIFI:STAT?").split(",")
+        if len(fields) != 4:
+            raise ScpiError(f"Unexpected Wi-Fi status response: {fields!r}")
+        return bool(int(fields[0])), int(fields[1]), int(fields[2]), int(fields[3])
+
+    def _write_line(self, command: str) -> None:
+        self.transport.write(command.encode("ascii") + b"\n")
+        self.transport.flush()
+
+    def _read_exact(self, size: int) -> bytes:
+        data = self.transport.read(size)
+        if len(data) != size:
+            raise ScpiTimeoutError(f"Expected {size} bytes, received {len(data)}.")
+        return data
+
+    def _drain_block_terminator(self) -> None:
+        try:
+            char = self.transport.read(1)
+        except TimeoutError:
+            return
+        if char in (b"\n", b"\r"):
+            return
+        # The firmware emits a newline after blocks. If a different byte appears,
+        # leave higher-level code to catch the protocol mismatch on the next read.
diff --git a/tests/test_pico_scpi_transport.py b/tests/test_pico_scpi_transport.py
new file mode 100644
index 0000000..f4e433f
--- /dev/null
+++ b/tests/test_pico_scpi_transport.py
@@ -0,0 +1,126 @@
+import pytest
+
+from pslab.pico import PicoDevice, ScpiClient, ScpiError
+
+
+class FakeTransport:
+    def __init__(self):
+        self.output = bytearray()
+        self.writes = []
+        self.timeout = 1.0
+        self.closed = False
+        self.connected = False
+
+    def connect(self):
+        self.connected = True
+
+    def close(self):
+        self.closed = True
+
+    def read(self, size):
+        data = self.output[:size]
+        del self.output[:size]
+        return bytes(data)
+
+    def readline(self):
+        try:
+            end = self.output.index(ord("\n")) + 1
+        except ValueError:
+            end = len(self.output)
+        return self.read(end)
+
+    def write(self, data):
+        self.writes.append(data)
+        command = data.strip().decode("ascii")
+        responses = {
+            "*IDN?": b"FOSSASIA,PSLab Pico,1.0,v0.1.0\n",
+            "SYST:ERR:COUN?": b"1\n",
+            "SYST:ERR?": b"-113,\"Undefined header\"\n",
+            "COMM:TRAN?": b"USB\n",
+            "COMM:WIFI:STAT?": b"1,12,3,4\n",
+        }
+        if command == "LA:READ?":
+            payload = b"abcd"
+            self.output.extend(b"#14" + payload + b"\n")
+        elif command == "BAD:BLOCK?":
+            self.output.extend(b"-200,\"Execution error\"\n")
+        elif command in responses:
+            self.output.extend(responses[command])
+        return len(data)
+
+    def flush(self):
+        pass
+
+    def reset_input_buffer(self):
+        self.output.clear()
+
+
+def test_query_writes_newline_and_reads_text_response():
+    transport = FakeTransport()
+    client = ScpiClient(transport)
+
+    response = client.query("*IDN?")
+
+    assert response == "FOSSASIA,PSLab Pico,1.0,v0.1.0"
+    assert transport.writes == [b"*IDN?\n"]
+
+
+def test_command_rejects_query_commands():
+    client = ScpiClient(FakeTransport())
+
+    with pytest.raises(ValueError):
+        client.command("*IDN?")
+
+
+def test_command_sends_silent_command():
+    transport = FakeTransport()
+    client = ScpiClient(transport)
+
+    client.command("*RST")
+
+    assert transport.writes == [b"*RST\n"]
+
+
+def test_query_block_reads_definite_length_payload():
+    client = ScpiClient(FakeTransport())
+
+    payload = client.query_block("LA:READ?")
+
+    assert payload == b"abcd"
+
+
+def test_query_block_raises_scpi_error_when_response_is_text_error():
+    client = ScpiClient(FakeTransport())
+
+    with pytest.raises(ScpiError, match="Execution error"):
+        client.query_block("BAD:BLOCK?")
+
+
+def test_system_error_parses_code_and_message():
+    client = ScpiClient(FakeTransport())
+
+    code, message = client.system_error()
+
+    assert code == -113
+    assert message == "Undefined header"
+
+
+def test_transport_helpers_parse_status_values():
+    transport = FakeTransport()
+    client = ScpiClient(transport)
+
+    client.set_transport("wifi")
+    assert client.get_transport() == "USB"
+    assert client.wifi_status() == (True, 12, 3, 4)
+    assert b"COMM:TRAN WIFI\n" in transport.writes
+
+
+def test_pico_device_wraps_scpi_client():
+    transport = FakeTransport()
+    device = PicoDevice.from_transport(transport)
+
+    with device:
+        assert device.identify() == "FOSSASIA,PSLab Pico,1.0,v0.1.0"
+
+    assert transport.connected
+    assert transport.closed