irtsp 0.1.0__py3-none-any.whl

irtsp/__init__.py

@@ -0,0 +1,135 @@
+"""irtsp — friendly Python client for iRTSP phones.
+
+Read your iPhone's camera + IMU / GPS / LiDAR-depth / ARKit-pose streams as
+typed, unit-aware (SI) records — all on one shared clock, so video and
+odometry align with **no time offset to estimate**.
+
+The 3-line tour::
+
+    import irtsp
+
+    with irtsp.connect("192.168.1.24") as phone:
+        for imu in phone.imu:
+            print(imu.gyro, imu.accel)     # rad/s, m/s²
+
+Pattern-match the whole odometry channel::
+
+    with irtsp.connect(phone_ip, depth=True) as phone:
+        for rec in phone.odometry:
+            match rec:
+                case irtsp.IMU(gyro=g):            ...
+                case irtsp.GNSS() as fix:          ...
+                case irtsp.Pose(position=p):       ...
+                case irtsp.DepthFrame() as d:      ...
+
+Extras: ``irtsp[numpy]`` (depth arrays), ``irtsp[discovery]`` (Bonjour),
+``irtsp[video]`` (frames + ``synced()`` bundles, experimental).
+
+Wire format & synchronization deep-dive:
+https://github.com/ryanrudes/irtsp-support/blob/main/INTEGRATION.md
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .clock import StreamClock, interpolate_pose, slerp
+from .records import (
+    STANDARD_GRAVITY,
+    Altitude,
+    DepthFrame,
+    GNSS,
+    Heading,
+    IMU,
+    Intrinsics,
+    Pose,
+    Quat,
+    RawAccel,
+    RawGyro,
+    Record,
+    Tracking,
+    Unknown,
+    Vec3,
+)
+from .session import Handshake, RecordStream, Session, connect
+from .wire import ConnectionClosed, ProtocolError, RecordType
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .aio import AsyncRecordStream, AsyncSession
+    from .discovery import Device
+
+__version__ = "0.1.0"
+
+# Names resolved lazily (PEP 562) so `import irtsp` stays dependency-free and
+# fast, while `irtsp.Device` / `irtsp.AsyncSession` still work at runtime.
+_LAZY = {
+    "Device": ("irtsp.discovery", "Device"),
+    "AsyncSession": ("irtsp.aio", "AsyncSession"),
+    "AsyncRecordStream": ("irtsp.aio", "AsyncRecordStream"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY:
+        import importlib
+
+        module_name, attr = _LAZY[name]
+        return getattr(importlib.import_module(module_name), attr)
+    raise AttributeError(f"module 'irtsp' has no attribute {name!r}")
+
+
+def __dir__() -> "list[str]":
+    return sorted(list(globals()) + list(_LAZY))
+
+__all__ = [
+    # connecting
+    "connect",
+    "connect_async",
+    "discover",
+    "Session",
+    "AsyncSession",
+    "AsyncRecordStream",
+    "Device",
+    "Handshake",
+    "RecordStream",
+    # records
+    "Record",
+    "IMU",
+    "RawGyro",
+    "RawAccel",
+    "Intrinsics",
+    "GNSS",
+    "Altitude",
+    "Heading",
+    "Pose",
+    "DepthFrame",
+    "Unknown",
+    # value types & constants
+    "Vec3",
+    "Quat",
+    "Tracking",
+    "STANDARD_GRAVITY",
+    # clock
+    "StreamClock",
+    "slerp",
+    "interpolate_pose",
+    # protocol
+    "RecordType",
+    "ProtocolError",
+    "ConnectionClosed",
+    "__version__",
+]
+
+
+def discover(timeout: float = 2.0) -> "list[Device]":
+    """Find iRTSP phones on the local network (needs ``irtsp[discovery]``)."""
+    from .discovery import discover as _discover
+
+    return _discover(timeout=timeout)
+
+
+async def connect_async(target: "str | Any", **kwargs: Any) -> "AsyncSession":
+    """Async twin of :func:`connect` — see :mod:`irtsp.aio`."""
+    from .aio import connect_async as _connect_async
+
+    return await _connect_async(target, **kwargs)

irtsp/aio.py

@@ -0,0 +1,355 @@
+"""asyncio flavor of the client — same shapes, ``async`` everywhere.
+
+::
+
+    import asyncio, irtsp
+
+    async def main():
+        async with await irtsp.connect_async("192.168.1.24") as phone:
+            async for imu in phone.imu:
+                print(imu.gyro, imu.accel)
+
+    asyncio.run(main())
+
+This is a native asyncio implementation (not a thread wrapper): one reader task
+per channel, fanning records out to independent bounded queues with the same
+drop-oldest-per-consumer policy as the sync :class:`~irtsp.Session`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import struct
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar
+
+from .clock import StreamClock
+from .records import (
+    Altitude,
+    DepthFrame,
+    GNSS,
+    Heading,
+    IMU,
+    Intrinsics,
+    Pose,
+    Record,
+)
+from .session import DEFAULT_BUFFER, Handshake, _check_handshake
+from .wire import (
+    MAX_MESSAGE,
+    RECORD_SIZE,
+    ProtocolError,
+    decode_depth_frame,
+    decode_record,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .discovery import Device
+
+__all__ = ["connect_async", "AsyncSession", "AsyncRecordStream"]
+
+log = logging.getLogger("irtsp.aio")
+
+R = TypeVar("R", bound=Record)
+
+_EOS: Any = object()  # end-of-stream sentinel
+
+
+async def _recv_handshake(reader: asyncio.StreamReader) -> Handshake:
+    (length,) = struct.unpack("<I", await reader.readexactly(4))
+    if not 0 < length <= MAX_MESSAGE:
+        raise ProtocolError(f"implausible handshake length {length}")
+    return Handshake(json.loads(await reader.readexactly(length)))
+
+
+class AsyncRecordStream(Generic[R]):
+    """Async-iterable, filtered view of the session's records."""
+
+    def __init__(self, session: "AsyncSession", types: tuple[type, ...], maxsize: int):
+        self._session = session
+        self.types = types
+        self._queue: asyncio.Queue[Any] = asyncio.Queue(maxsize=maxsize)
+        self._dropped = 0
+        self._ended = False
+
+    @property
+    def dropped(self) -> int:
+        """Records this consumer lost because it wasn't keeping up."""
+        return self._dropped
+
+    def _put(self, item: Any) -> None:
+        try:
+            self._queue.put_nowait(item)
+        except asyncio.QueueFull:
+            try:
+                self._queue.get_nowait()
+                self._dropped += 1
+            except asyncio.QueueEmpty:  # pragma: no cover — racy but harmless
+                pass
+            self._queue.put_nowait(item)
+
+    def pop_all(self) -> list[R]:
+        """Drain everything buffered right now, without blocking."""
+        out: list[R] = []
+        while True:
+            try:
+                item = self._queue.get_nowait()
+            except asyncio.QueueEmpty:
+                return out
+            if item is _EOS:
+                self._ended = True
+                self._queue.put_nowait(_EOS)  # keep the sentinel for __anext__
+                return out
+            out.append(item)
+
+    def close(self) -> None:
+        """Unsubscribe. Any pending ``async for`` ends."""
+        self._session._unsubscribe(self)
+        self._put(_EOS)
+
+    def __aiter__(self) -> "AsyncRecordStream[R]":
+        return self
+
+    async def __anext__(self) -> R:
+        if self._ended:
+            raise StopAsyncIteration
+        item = await self._queue.get()
+        if item is _EOS:
+            self._ended = True
+            raise StopAsyncIteration
+        return item
+
+
+class AsyncSession:
+    """A live async connection. Prefer :func:`irtsp.connect_async`."""
+
+    def __init__(self, host: str, *, imu_port: int = 8555, depth: bool = False,
+                 depth_port: int = 8556, buffer: int = DEFAULT_BUFFER):
+        self.host = host
+        self.imu_port = imu_port
+        self.depth_enabled = depth
+        self.depth_port = depth_port
+        self._buffer = buffer
+
+        self.info: Handshake | None = None
+        self.depth_info: Handshake | None = None
+
+        self._subs: list[AsyncRecordStream[Any]] = []
+        self._callbacks: list[tuple[tuple[type, ...], Callable[[Any], None]]] = []
+        self._latest: dict[type, Record] = {}
+        self._tasks: list[asyncio.Task[None]] = []
+        self._writers: list[asyncio.StreamWriter] = []
+        self._closed = False
+
+    # ------------------------------------------------------------------ setup
+
+    async def open(self) -> "AsyncSession":
+        try:
+            reader, writer = await asyncio.open_connection(self.host, self.imu_port)
+            self._writers.append(writer)
+            info = await _recv_handshake(reader)
+            _check_handshake(info, "irtsp-imu")
+            self.info = info
+            self._tasks.append(
+                asyncio.create_task(self._record_loop(reader), name="irtsp-odometry")
+            )
+
+            if self.depth_enabled:
+                dreader, dwriter = await asyncio.open_connection(self.host, self.depth_port)
+                self._writers.append(dwriter)
+                depth_info = await _recv_handshake(dreader)
+                _check_handshake(depth_info, "irtsp-depth")
+                self.depth_info = depth_info
+                self._tasks.append(
+                    asyncio.create_task(self._depth_loop(dreader), name="irtsp-depth")
+                )
+        except BaseException:
+            await self.aclose()  # tear down whatever already started
+            raise
+        return self
+
+    # ------------------------------------------------------------ reader tasks
+
+    async def _record_loop(self, reader: asyncio.StreamReader) -> None:
+        last_seq: int | None = None
+        try:
+            while True:
+                buf = await reader.readexactly(RECORD_SIZE)
+                record = decode_record(buf)
+                # Server replays the latest Intrinsics (stale seq) to late
+                # joiners — don't let it baseline the gap tracker.
+                if last_seq is None and isinstance(record, Intrinsics):
+                    self._dispatch(record, None)
+                    continue
+                last_seq = self._dispatch(record, last_seq)
+        except (asyncio.IncompleteReadError, ConnectionError, OSError):
+            pass
+        except ProtocolError as e:
+            log.error("odometry stream desynced: %s", e)
+        finally:
+            await self.aclose()
+
+    async def _depth_loop(self, reader: asyncio.StreamReader) -> None:
+        last_seq: int | None = None
+        try:
+            while True:
+                (length,) = struct.unpack("<I", await reader.readexactly(4))
+                if not 0 < length <= MAX_MESSAGE:
+                    raise ProtocolError(f"implausible depth frame length {length}")
+                frame = decode_depth_frame(await reader.readexactly(length))
+                last_seq = self._dispatch(frame, last_seq)
+        except (asyncio.IncompleteReadError, ConnectionError, OSError):
+            pass
+        except ProtocolError as e:
+            log.error("depth stream desynced: %s", e)
+        finally:
+            await self.aclose()
+
+    def _dispatch(self, record: Record, last_seq: int | None) -> int:
+        if last_seq is not None:
+            gap = (record.seq - last_seq - 1) & 0xFFFF
+            if gap:
+                record = replace(record, gap=gap)
+        self._latest[type(record)] = record
+        for sub in self._subs:
+            if isinstance(record, sub.types):
+                sub._put(record)
+        for types, fn in self._callbacks:
+            if isinstance(record, types):
+                try:
+                    fn(record)
+                except Exception:
+                    log.exception("irtsp callback %r raised", fn)
+        return record.seq
+
+    # -------------------------------------------------------------- consuming
+
+    def stream(self, *types: type[R], buffer: int | None = None) -> AsyncRecordStream[R]:
+        """A new independent async iterator over the given record types."""
+        sub: AsyncRecordStream[R] = AsyncRecordStream(
+            self, types or (Record,), buffer or self._buffer
+        )
+        if self._closed:
+            sub._put(_EOS)
+        else:
+            self._subs.append(sub)
+        return sub
+
+    def _unsubscribe(self, sub: AsyncRecordStream[Any]) -> None:
+        if sub in self._subs:
+            self._subs.remove(sub)
+
+    @property
+    def odometry(self) -> AsyncRecordStream[Record]:
+        return self.stream()
+
+    @property
+    def imu(self) -> AsyncRecordStream[IMU]:
+        return self.stream(IMU)
+
+    @property
+    def intrinsics(self) -> AsyncRecordStream[Intrinsics]:
+        return self.stream(Intrinsics)
+
+    @property
+    def gnss(self) -> AsyncRecordStream[GNSS]:
+        return self.stream(GNSS)
+
+    @property
+    def altitude(self) -> AsyncRecordStream[Altitude]:
+        return self.stream(Altitude)
+
+    @property
+    def heading(self) -> AsyncRecordStream[Heading]:
+        return self.stream(Heading)
+
+    @property
+    def pose(self) -> AsyncRecordStream[Pose]:
+        return self.stream(Pose)
+
+    @property
+    def depth(self) -> AsyncRecordStream[DepthFrame]:
+        if not self.depth_enabled:
+            raise RuntimeError("depth channel not enabled — connect_async(host, depth=True)")
+        return self.stream(DepthFrame)
+
+    def on(self, types: type[R] | tuple[type[R], ...], callback: Callable[[R], None]) -> None:
+        """Call ``callback(record)`` for matching records (on the event loop)."""
+        filt = types if isinstance(types, tuple) else (types,)
+        self._callbacks.append((filt, callback))
+
+    def latest(self, rtype: type[R]) -> R | None:
+        """The most recent record of exactly ``rtype`` seen, or ``None``."""
+        return self._latest.get(rtype)  # type: ignore[return-value]
+
+    @property
+    def clock(self) -> StreamClock:
+        if self.info is None:
+            raise RuntimeError("session is not open")
+        return self.info.clock
+
+    # ---------------------------------------------------------------- teardown
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    async def aclose(self) -> None:
+        """Close connections and end all iterators. Idempotent."""
+        if self._closed:
+            return
+        self._closed = True
+        for writer in self._writers:
+            try:
+                writer.close()
+            except Exception:  # pragma: no cover
+                pass
+        for sub in self._subs:
+            sub._put(_EOS)
+        # Cancel and await the reader tasks — except the current one: the
+        # loops call aclose() from their own finally, and awaiting yourself
+        # deadlocks.
+        current = asyncio.current_task()
+        others = [t for t in self._tasks if t is not current]
+        for task in others:
+            task.cancel()
+        if others:
+            await asyncio.gather(*others, return_exceptions=True)
+        for writer in self._writers:
+            try:
+                await writer.wait_closed()
+            except Exception:  # pragma: no cover — transport already gone
+                pass
+
+    async def __aenter__(self) -> "AsyncSession":
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.aclose()
+
+    def __repr__(self) -> str:  # pragma: no cover
+        state = "closed" if self._closed else "live"
+        return f"<irtsp.AsyncSession {self.host}:{self.imu_port} [{state}]>"
+
+
+async def connect_async(
+    target: "str | Device",
+    *,
+    imu_port: int = 8555,
+    depth: bool = False,
+    depth_port: int = 8556,
+    buffer: int = DEFAULT_BUFFER,
+) -> AsyncSession:
+    """Async twin of :func:`irtsp.connect` (odometry + depth channels)."""
+    host = target
+    if not isinstance(target, str):
+        host = target.host
+        ports = getattr(target, "ports", {})
+        imu_port = ports.get("imu", imu_port)
+        depth_port = ports.get("depth", depth_port)
+    return await AsyncSession(
+        host,  # type: ignore[arg-type]
+        imu_port=imu_port, depth=depth, depth_port=depth_port, buffer=buffer
+    ).open()

irtsp/clock.py

@@ -0,0 +1,99 @@
+"""The shared clock, and time interpolation helpers.
+
+iRTSP captures **one** anchor pair per streaming session — a host-clock reading
+and the Unix wall time at the same instant — and every stream (video RTP/RTCP,
+IMU, GPS, pose, depth) derives its timestamps from it. That is why the streams
+need no cross-correlation to align: they were never on different clocks.
+
+:class:`StreamClock` is that anchor, parsed straight out of the handshake, and
+converts between the two axes. See the integration guide §3–4 for the full
+story of how the video's RTCP Sender Reports land on the same ``unix_ts`` axis.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+from .records import Pose, Quat, Tracking, Vec3
+
+__all__ = ["StreamClock", "slerp", "interpolate_pose"]
+
+
+@dataclass(frozen=True, slots=True)
+class StreamClock:
+    """The session's host↔wall anchor: ``unix = wall_anchor + (host − host_anchor)``."""
+
+    host_anchor: float  #: host-clock seconds (mach_absolute_time) at anchor instant
+    wall_anchor: float  #: Unix wall seconds at the same instant
+
+    @classmethod
+    def from_handshake(cls, handshake: Mapping[str, Any]) -> "StreamClock":
+        clock = handshake.get("clock", {})
+        return cls(
+            host_anchor=float(clock.get("host_anchor", 0.0)),
+            wall_anchor=float(clock.get("wall_anchor", 0.0)),
+        )
+
+    def to_unix(self, host_ts: float) -> float:
+        """Map a host-clock timestamp onto the wall clock (the RTCP-SR NTP axis)."""
+        return self.wall_anchor + (host_ts - self.host_anchor)
+
+    def to_host(self, unix_ts: float) -> float:
+        """Map a wall-clock timestamp back onto the host clock (the video-PTS axis)."""
+        return self.host_anchor + (unix_ts - self.wall_anchor)
+
+
+def slerp(a: Quat, b: Quat, t: float) -> Quat:
+    """Spherical linear interpolation between two unit quaternions, ``t ∈ [0, 1]``.
+
+    Takes the shortest path, and degrades gracefully to normalized lerp when the
+    quaternions are nearly parallel.
+    """
+    ax, ay, az, aw = a
+    bx, by, bz, bw = b
+    dot = ax * bx + ay * by + az * bz + aw * bw
+    if dot < 0.0:  # shortest path
+        bx, by, bz, bw, dot = -bx, -by, -bz, -bw, -dot
+
+    if dot > 0.9995:  # nearly parallel — nlerp avoids division blow-up
+        q = Quat(
+            ax + t * (bx - ax),
+            ay + t * (by - ay),
+            az + t * (bz - az),
+            aw + t * (bw - aw),
+        )
+        return q.normalized()
+
+    theta0 = math.acos(min(dot, 1.0))
+    theta = theta0 * t
+    sin_theta0 = math.sin(theta0)
+    s0 = math.sin(theta0 - theta) / sin_theta0
+    s1 = math.sin(theta) / sin_theta0
+    return Quat(ax * s0 + bx * s1, ay * s0 + by * s1, az * s0 + bz * s1, aw * s0 + bw * s1)
+
+
+def interpolate_pose(a: Pose, b: Pose, unix_ts: float) -> Pose:
+    """The pose at ``unix_ts``, interpolated between two bracketing samples.
+
+    Position is interpolated linearly and orientation spherically; ``tracking``
+    reports the *worse* of the two endpoints so degraded tracking is never
+    hidden by interpolation. ``unix_ts`` outside ``[a, b]`` is clamped.
+    """
+    if b.unix_ts <= a.unix_ts:
+        return a
+    t = max(0.0, min(1.0, (unix_ts - a.unix_ts) / (b.unix_ts - a.unix_ts)))
+    position = Vec3(
+        a.position.x + t * (b.position.x - a.position.x),
+        a.position.y + t * (b.position.y - a.position.y),
+        a.position.z + t * (b.position.z - a.position.z),
+    )
+    return Pose(
+        host_ts=a.host_ts + t * (b.host_ts - a.host_ts),
+        unix_ts=a.unix_ts + t * (b.unix_ts - a.unix_ts),
+        seq=b.seq,
+        position=position,
+        orientation=slerp(a.orientation, b.orientation, t),
+        tracking=Tracking(min(a.tracking, b.tracking)),
+    )

irtsp/discovery.py

@@ -0,0 +1,113 @@
+"""Find iRTSP phones on the local network via Bonjour/mDNS.
+
+Requires the ``discovery`` extra::
+
+    pip install 'irtsp[discovery]'
+
+Then::
+
+    import irtsp
+
+    for device in irtsp.discover():
+        print(device.name, device.host, device.ports)
+
+    phone = irtsp.discover()[0].connect()
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .session import Session
+
+__all__ = ["Device", "discover"]
+
+#: Bonjour service type → channel name.
+_SERVICE_CHANNELS = {
+    "_rtsp._tcp.local.": "video",
+    "_irtsp-imu._tcp.local.": "imu",
+    "_irtsp-depth._tcp.local.": "depth",
+}
+
+
+@dataclass
+class Device:
+    """One discovered iRTSP phone."""
+
+    name: str  #: the service name the phone advertises
+    host: str  #: best address to connect to (IPv4 preferred)
+    addresses: list[str] = field(default_factory=list)
+    ports: dict[str, int] = field(default_factory=dict)  #: e.g. ``{"video": 8554, "imu": 8555}``
+    properties: dict[str, str] = field(default_factory=dict)
+
+    def connect(self, **kwargs: Any) -> "Session":
+        """Open a session to this phone — same options as :func:`irtsp.connect`."""
+        from .session import connect
+
+        return connect(self, **kwargs)
+
+    def __repr__(self) -> str:  # pragma: no cover
+        ports = ", ".join(f"{k}:{v}" for k, v in sorted(self.ports.items()))
+        return f"<irtsp.Device {self.name!r} {self.host} ({ports})>"
+
+
+def discover(timeout: float = 2.0) -> list[Device]:
+    """Browse the local network for iRTSP phones for ``timeout`` seconds.
+
+    Only devices advertising the iRTSP odometry channel (``_irtsp-imu._tcp``)
+    are returned — a plain RTSP camera that isn't iRTSP won't show up.
+    """
+    try:
+        from zeroconf import ServiceBrowser, ServiceListener, Zeroconf
+    except ImportError as e:
+        raise ImportError(
+            "irtsp.discover() needs zeroconf — pip install 'irtsp[discovery]' "
+            "(or connect by IP: irtsp.connect('192.168.1.24'))"
+        ) from e
+
+    found: dict[str, Device] = {}
+
+    class _Listener(ServiceListener):
+        def add_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+            info = zc.get_service_info(type_, name, timeout=int(timeout * 1000))
+            if info is None:
+                return
+            service_name = name.removesuffix("." + type_)
+            addresses = [addr for addr in info.parsed_scoped_addresses() if ":" not in addr] or list(
+                info.parsed_scoped_addresses()
+            )
+            if not addresses:
+                return
+            device = found.setdefault(
+                service_name, Device(name=service_name, host=addresses[0], addresses=addresses)
+            )
+            device.ports[_SERVICE_CHANNELS[type_]] = info.port or 0
+            for k, v in (info.properties or {}).items():
+                try:
+                    device.properties[k.decode()] = v.decode() if isinstance(v, bytes) else str(v)
+                except Exception:  # pragma: no cover — malformed TXT records
+                    pass
+
+        def update_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+            self.add_service(zc, type_, name)
+
+        def remove_service(self, zc: Zeroconf, type_: str, name: str) -> None:
+            pass
+
+    zc = Zeroconf()
+    try:
+        listener = _Listener()
+        browsers = [ServiceBrowser(zc, t, listener) for t in _SERVICE_CHANNELS]
+        time.sleep(timeout)
+        del browsers
+    finally:
+        zc.close()
+
+    # A phone must expose the odometry channel to count as an iRTSP device.
+    return sorted(
+        (d for d in found.values() if "imu" in d.ports),
+        key=lambda d: d.name,
+    )