jmri_loconet.messages

  1"""Typed LocoNet message dataclasses + encode/decode dispatch.
  2
  3Each opcode JMRI's ``LnConstants`` declares maps to one frozen dataclass.
  4:func:`decode` turns frame bytes into the matching dataclass;
  5:func:`encode` does the reverse (returns frame bytes including checksum).
  6
  7Unrecognized but well-formed frames decode to :class:`OpcUnknown` so
  8listeners can still observe them — JMRI exposes everything on the bus.
  9"""
 10from __future__ import annotations
 11
 12from dataclasses import dataclass
 13
 14from jmri_loconet.codec import append_checksum, compute_checksum, verify_checksum
 15from jmri_loconet.opcodes import (
 16    OPC_GPOFF,
 17    OPC_GPON,
 18    OPC_IDLE,
 19    OPC_INPUT_REP,
 20    OPC_LINK_SLOTS,
 21    OPC_LOCO_ADR,
 22    OPC_LOCO_DIRF,
 23    OPC_LOCO_SND,
 24    OPC_LOCO_SPD,
 25    OPC_LONG_ACK,
 26    OPC_MOVE_SLOTS,
 27    OPC_RQ_SL_DATA,
 28    OPC_SL_RD_DATA,
 29    OPC_SLOT_STAT1,
 30    OPC_SW_REP,
 31    OPC_SW_REQ,
 32    OPC_UNLINK_SLOTS,
 33    OPC_WR_SL_DATA,
 34)
 35
 36# ---------- base + unknown ---------------------------------------------------
 37
 38
 39class LocoNetMessage:
 40    """Marker base type. Subclasses are frozen dataclasses."""
 41    opcode: int
 42
 43
 44@dataclass(frozen=True, slots=True)
 45class OpcUnknown(LocoNetMessage):
 46    """A well-formed frame whose opcode we don't have a typed wrapper for."""
 47    opcode: int
 48    payload: bytes  # everything between opcode and checksum (exclusive)
 49
 50
 51# ---------- 2-byte messages --------------------------------------------------
 52
 53
 54@dataclass(frozen=True, slots=True)
 55class OpcGpOn(LocoNetMessage):
 56    opcode: int = OPC_GPON
 57
 58
 59@dataclass(frozen=True, slots=True)
 60class OpcGpOff(LocoNetMessage):
 61    opcode: int = OPC_GPOFF
 62
 63
 64@dataclass(frozen=True, slots=True)
 65class OpcIdle(LocoNetMessage):
 66    opcode: int = OPC_IDLE
 67
 68
 69# ---------- 4-byte messages --------------------------------------------------
 70
 71
 72@dataclass(frozen=True, slots=True)
 73class OpcSwReq(LocoNetMessage):
 74    """Switch request — throw or close a turnout.
 75
 76    Address is 1..2048 (LocoNet packs it across two bytes with bit layout
 77    that JMRI documents explicitly).
 78    """
 79    address: int
 80    closed: bool       # True = closed direction, False = thrown
 81    output_on: bool    # True = energize, False = de-energize
 82    opcode: int = OPC_SW_REQ
 83
 84
 85@dataclass(frozen=True, slots=True)
 86class OpcSwRep(LocoNetMessage):
 87    """Switch report — turnout state feedback from the layout."""
 88    address: int
 89    closed: bool
 90    thrown: bool
 91    opcode: int = OPC_SW_REP
 92
 93
 94@dataclass(frozen=True, slots=True)
 95class OpcInputRep(LocoNetMessage):
 96    """Sensor input report."""
 97    address: int
 98    active: bool
 99    opcode: int = OPC_INPUT_REP
100
101
102@dataclass(frozen=True, slots=True)
103class OpcLongAck(LocoNetMessage):
104    """Long-form acknowledgment — sent in response to many request opcodes."""
105    responding_opcode: int
106    ack_code: int
107    opcode: int = OPC_LONG_ACK
108
109
110@dataclass(frozen=True, slots=True)
111class OpcSlotStat1(LocoNetMessage):
112    slot: int
113    stat: int
114    opcode: int = OPC_SLOT_STAT1
115
116
117@dataclass(frozen=True, slots=True)
118class OpcMoveSlots(LocoNetMessage):
119    """Move slot src → dest. Same src+dest=null move (refresh); 0=dispatch."""
120    src: int
121    dest: int
122    opcode: int = OPC_MOVE_SLOTS
123
124
125@dataclass(frozen=True, slots=True)
126class OpcLinkSlots(LocoNetMessage):
127    """Link ``src`` as a follower of ``dest`` — build a consist chain.
128
129    After a LINK, the command station broadcasts speed and direction
130    commands from *dest* to *src* automatically; the client only needs
131    to drive the lead loco. Functions remain per-loco.
132    """
133    src: int
134    dest: int
135    opcode: int = OPC_LINK_SLOTS
136
137
138@dataclass(frozen=True, slots=True)
139class OpcUnlinkSlots(LocoNetMessage):
140    """Break the LINK between *src* and *dest*.
141
142    Ordering: *dest* is the slot higher in the chain, *src* is the
143    follower being detached. After UNLINK, *src* is a standalone slot
144    again and stops receiving speed/direction fan-out from *dest*.
145    """
146    src: int
147    dest: int
148    opcode: int = OPC_UNLINK_SLOTS
149
150
151@dataclass(frozen=True, slots=True)
152class OpcRqSlData(LocoNetMessage):
153    """Request slot data — server replies with OpcSlRdData."""
154    slot: int
155    opcode: int = OPC_RQ_SL_DATA
156
157
158@dataclass(frozen=True, slots=True)
159class OpcLocoAdr(LocoNetMessage):
160    """Request the slot currently holding *address*.
161
162    Long addresses use the high byte; short addresses set high=0.
163    """
164    address: int
165    opcode: int = OPC_LOCO_ADR
166
167
168@dataclass(frozen=True, slots=True)
169class OpcLocoSpd(LocoNetMessage):
170    """Set speed for a slot. ``speed`` is 0..127."""
171    slot: int
172    speed: int
173    opcode: int = OPC_LOCO_SPD
174
175
176@dataclass(frozen=True, slots=True)
177class OpcLocoDirF(LocoNetMessage):
178    """Set direction + F0..F4 for a slot.
179
180    ``dirf`` is the wire byte: bit 5 = direction (0 = forward), bits 0..4 = F0..F4.
181    """
182    slot: int
183    dirf: int
184    opcode: int = OPC_LOCO_DIRF
185
186
187@dataclass(frozen=True, slots=True)
188class OpcLocoSnd(LocoNetMessage):
189    """Set F5..F8 for a slot. ``snd`` packs F5..F8 in bits 0..3."""
190    slot: int
191    snd: int
192    opcode: int = OPC_LOCO_SND
193
194
195# ---------- variable-length messages -----------------------------------------
196
197
198@dataclass(frozen=True, slots=True)
199class OpcSlRdData(LocoNetMessage):
200    """Slot read data — 14-byte response carrying full slot state.
201
202    ``ss2`` (byte index 8 on the wire) holds the "STAT2" byte for
203    regular slots, and the encoded hours field for the fast-clock
204    slot (123). It defaults to 0 so existing callers don't need to
205    care; the clock module fills it in.
206    """
207    slot: int
208    stat: int
209    address: int        # composed from ADR + ADR2
210    speed: int
211    dirf: int
212    trk: int
213    snd: int
214    ss2: int = 0
215    opcode: int = OPC_SL_RD_DATA
216
217
218@dataclass(frozen=True, slots=True)
219class OpcWrSlData(LocoNetMessage):
220    """Write slot data — same shape as OpcSlRdData but client → master."""
221    slot: int
222    stat: int
223    address: int
224    speed: int
225    dirf: int
226    trk: int
227    snd: int
228    ss2: int = 0
229    opcode: int = OPC_WR_SL_DATA
230
231
232# ---------- decode -----------------------------------------------------------
233
234
235def decode(frame: bytes) -> LocoNetMessage:
236    """Turn a complete, checksum-verified frame into a typed message.
237
238    Raises :class:`ValueError` if the checksum is invalid.
239    """
240    if not verify_checksum(frame):
241        raise ValueError("LocoNet frame has invalid checksum")
242    op = frame[0]
243
244    # 2-byte
245    if op == OPC_GPON:
246        return OpcGpOn()
247    if op == OPC_GPOFF:
248        return OpcGpOff()
249    if op == OPC_IDLE:
250        return OpcIdle()
251
252    # 4-byte: layout is [op, b1, b2, cs]; b1/b2 vary per opcode.
253    if op == OPC_SW_REQ:
254        return OpcSwReq(
255            address=_decode_switch_address(frame[1], frame[2]),
256            closed=bool(frame[2] & 0x20),
257            output_on=bool(frame[2] & 0x10),
258        )
259    if op == OPC_SW_REP:
260        return OpcSwRep(
261            address=_decode_switch_address(frame[1], frame[2]),
262            closed=bool(frame[2] & 0x20),
263            thrown=not bool(frame[2] & 0x20),
264        )
265    if op == OPC_INPUT_REP:
266        # JMRI's mapping:
267        #   pair = ((SN2 & 0x0F) << 7) | (SN1 & 0x7F)
268        #   sensor (1-based) = pair * 2 + ((SN2 & 0x20) >> 5) + 1
269        pair = ((frame[2] & 0x0F) << 7) | (frame[1] & 0x7F)
270        sensor_addr = (pair * 2) + ((frame[2] & 0x20) >> 5) + 1
271        active = bool(frame[2] & 0x10)
272        return OpcInputRep(address=sensor_addr, active=active)
273    if op == OPC_LONG_ACK:
274        return OpcLongAck(responding_opcode=frame[1] | 0x80, ack_code=frame[2])
275    if op == OPC_SLOT_STAT1:
276        return OpcSlotStat1(slot=frame[1], stat=frame[2])
277    if op == OPC_MOVE_SLOTS:
278        return OpcMoveSlots(src=frame[1], dest=frame[2])
279    if op == OPC_LINK_SLOTS:
280        return OpcLinkSlots(src=frame[1], dest=frame[2])
281    if op == OPC_UNLINK_SLOTS:
282        return OpcUnlinkSlots(src=frame[1], dest=frame[2])
283    if op == OPC_RQ_SL_DATA:
284        return OpcRqSlData(slot=frame[1])
285    if op == OPC_LOCO_ADR:
286        return OpcLocoAdr(address=(frame[1] << 7) | frame[2])
287    if op == OPC_LOCO_SPD:
288        return OpcLocoSpd(slot=frame[1], speed=frame[2])
289    if op == OPC_LOCO_DIRF:
290        return OpcLocoDirF(slot=frame[1], dirf=frame[2])
291    if op == OPC_LOCO_SND:
292        return OpcLocoSnd(slot=frame[1], snd=frame[2])
293
294    # Variable-length: 14-byte slot data frames.
295    if op == OPC_SL_RD_DATA and len(frame) == 14:
296        return OpcSlRdData(
297            slot=frame[2], stat=frame[3],
298            address=(frame[9] << 7) | frame[4],
299            speed=frame[5], dirf=frame[6], trk=frame[7],
300            ss2=frame[8], snd=frame[10],
301        )
302    if op == OPC_WR_SL_DATA and len(frame) == 14:
303        return OpcWrSlData(
304            slot=frame[2], stat=frame[3],
305            address=(frame[9] << 7) | frame[4],
306            speed=frame[5], dirf=frame[6], trk=frame[7],
307            ss2=frame[8], snd=frame[10],
308        )
309
310    return OpcUnknown(opcode=op, payload=bytes(frame[1:-1]))
311
312
313# ---------- encode -----------------------------------------------------------
314
315
316def encode(msg: LocoNetMessage) -> bytes:
317    """Turn *msg* into a complete frame (opcode + payload + checksum)."""
318    op = msg.opcode
319
320    # 2-byte
321    if op in (OPC_GPON, OPC_GPOFF, OPC_IDLE):
322        return append_checksum(bytes([op]))
323
324    # 4-byte — isinstance narrows the type cleanly per branch.
325    if isinstance(msg, OpcSwReq):
326        b1, b2 = _encode_switch_address(
327            msg.address, closed=msg.closed, output_on=msg.output_on
328        )
329        return append_checksum(bytes([op, b1, b2]))
330    if isinstance(msg, OpcSwRep):
331        b1, b2 = _encode_switch_address(
332            msg.address, closed=msg.closed, output_on=False, report_mode=True
333        )
334        return append_checksum(bytes([op, b1, b2]))
335    if isinstance(msg, OpcInputRep):
336        addr = msg.address - 1
337        sensor_bit = addr & 0x01
338        addr_words = addr >> 1
339        b1 = addr_words & 0x7F
340        b2 = ((addr_words >> 7) & 0x0F) | (sensor_bit << 5) | (0x10 if msg.active else 0)
341        return append_checksum(bytes([op, b1, b2]))
342    if isinstance(msg, OpcLongAck):
343        return append_checksum(
344            bytes([op, msg.responding_opcode & 0x7F, msg.ack_code & 0x7F])
345        )
346    if isinstance(msg, OpcSlotStat1):
347        return append_checksum(bytes([op, msg.slot & 0x7F, msg.stat & 0x7F]))
348    if isinstance(msg, OpcMoveSlots):
349        return append_checksum(bytes([op, msg.src & 0x7F, msg.dest & 0x7F]))
350    if isinstance(msg, OpcLinkSlots):
351        return append_checksum(bytes([op, msg.src & 0x7F, msg.dest & 0x7F]))
352    if isinstance(msg, OpcUnlinkSlots):
353        return append_checksum(bytes([op, msg.src & 0x7F, msg.dest & 0x7F]))
354    if isinstance(msg, OpcRqSlData):
355        return append_checksum(bytes([op, msg.slot & 0x7F, 0x00]))
356    if isinstance(msg, OpcLocoAdr):
357        return append_checksum(
358            bytes([op, (msg.address >> 7) & 0x7F, msg.address & 0x7F])
359        )
360    if isinstance(msg, OpcLocoSpd):
361        return append_checksum(bytes([op, msg.slot & 0x7F, msg.speed & 0x7F]))
362    if isinstance(msg, OpcLocoDirF):
363        return append_checksum(bytes([op, msg.slot & 0x7F, msg.dirf & 0x7F]))
364    if isinstance(msg, OpcLocoSnd):
365        return append_checksum(bytes([op, msg.slot & 0x7F, msg.snd & 0x7F]))
366
367    # 14-byte slot data frames.
368    if isinstance(msg, (OpcSlRdData, OpcWrSlData)):
369        body = bytes([
370            op,
371            14,                  # length byte
372            msg.slot & 0x7F,
373            msg.stat & 0x7F,
374            msg.address & 0x7F,    # ADR (low)
375            msg.speed & 0x7F,
376            msg.dirf & 0x7F,
377            msg.trk & 0x7F,
378            msg.ss2 & 0x7F,      # SS2 — hours for slot 123, STAT2 elsewhere
379            (msg.address >> 7) & 0x7F,  # ADR2 (high)
380            msg.snd & 0x7F,
381            0x00,                # ID1
382            0x00,                # ID2
383        ])
384        return append_checksum(body)
385
386    if isinstance(msg, OpcUnknown):
387        body = bytes([msg.opcode]) + msg.payload
388        return append_checksum(body)
389
390    raise NotImplementedError(f"encode not implemented for opcode 0x{op:02X}")
391
392
393# ---------- internals --------------------------------------------------------
394
395
396def _decode_switch_address(b1: int, b2: int) -> int:
397    """Decode a LocoNet switch (turnout) address from the two body bytes."""
398    # b1 = lower 7 bits, b2 low 4 bits = upper 4 bits; address is 1-based.
399    return ((b2 & 0x0F) << 7 | (b1 & 0x7F)) + 1
400
401
402def _encode_switch_address(
403    address: int, *, closed: bool, output_on: bool, report_mode: bool = False
404) -> tuple[int, int]:
405    """Encode a switch address into the two body bytes.
406
407    *report_mode* is True for OPC_SW_REP frames (bit-3 layout differs slightly
408    in JMRI; the visible difference is which bits carry feedback flags).
409    """
410    addr0 = address - 1
411    b1 = addr0 & 0x7F
412    b2 = (addr0 >> 7) & 0x0F
413    if closed:
414        b2 |= 0x20
415    if output_on and not report_mode:
416        b2 |= 0x10
417    return b1, b2
418
419
420# Re-export so the test suite can compute its own checksums.
421__all__ = [
422    "LocoNetMessage",
423    "OpcGpOff",
424    "OpcGpOn",
425    "OpcIdle",
426    "OpcInputRep",
427    "OpcLinkSlots",
428    "OpcLocoAdr",
429    "OpcLocoDirF",
430    "OpcLocoSnd",
431    "OpcLocoSpd",
432    "OpcLongAck",
433    "OpcMoveSlots",
434    "OpcUnlinkSlots",
435    "OpcRqSlData",
436    "OpcSlRdData",
437    "OpcSlotStat1",
438    "OpcSwRep",
439    "OpcSwReq",
440    "OpcUnknown",
441    "OpcWrSlData",
442    "compute_checksum",
443    "decode",
444    "encode",
445]
446
447# Re-export used internally so static analyzers don't flag the import.
448_ = compute_checksum