jmri_core.beans.consist_manager

  1"""``ConsistManager`` — CRUD for consists + fan-out throttle control.
  2
  3A consist is an ordered set of locomotives that operate as one — an MU
  4(multi-unit) group. Three flavours exist in the DCC world:
  5
  6* **Advanced** (aka *NMRA*, aka *CV19*) — each loco's CV19 holds a
  7  shared short address; the decoder itself responds to commands on
  8  that address. Set-up is a decoder programming task; running the
  9  group is transparent afterwards.
 10* **Command station** — the CS maintains the group in its own memory
 11  and forwards commands sent to the group's "handle" to every
 12  member. LocoNet's ``OPC_LINK_SLOTS``, NCE's consist memory table,
 13  and DCC-EX's ``<C>`` command are all this shape.
 14* **Software** — JMRI holds the group; every command sent to the
 15  group fans out into individual throttle commands, one per member.
 16  Works on any driver because it only uses the ``ThrottleManager``
 17  API — the fallback for drivers that don't (or can't) do hardware
 18  consisting.
 19
 20This module ships:
 21
 22* the base :class:`ConsistManager` — CRUD + fan-out helpers that work
 23  for every driver via the connection's :class:`ThrottleManager`;
 24* subclass hooks so drivers can override the wire ops without
 25  changing the public surface.
 26
 27Driver overrides live in each driver package (``jmri_loconet.consist``,
 28``jmri_nce.consist``, etc.) and are wired via
 29:meth:`SystemConnection._make_consist_manager`.
 30"""
 31from __future__ import annotations
 32
 33from collections.abc import Iterator
 34from typing import TYPE_CHECKING
 35
 36from jmri_core.beans.consist import Consist, ConsistKind, ConsistLoco
 37from jmri_core.events import Signal
 38from jmri_core.state import ThrottleDirection
 39
 40if TYPE_CHECKING:
 41    from jmri_core.beans.programmer import ProgrammerManager
 42    from jmri_core.beans.throttle_manager import ThrottleManager
 43
 44_log = __import__("logging").getLogger(__name__)
 45
 46
 47class ConsistManager:
 48    """Track a set of consists and fan throttle commands to their members.
 49
 50    The base implementation is the **software consist** manager: it
 51    stores consists in a dict keyed by the group's DCC address and
 52    executes group speed/direction/function commands by looping over
 53    every member's :class:`Throttle`. Drivers that support hardware
 54    consisting subclass this and override :meth:`_after_create` and
 55    :meth:`_before_delete` to install/remove the wire-level group;
 56    they typically leave the fan-out helpers alone since the base's
 57    per-loco commands are still valid even when the CS is grouping.
 58
 59    ``created`` and ``deleted`` are :class:`Signal` fan-outs so the
 60    push side of the JSON server, the WiThrottle broadcast, and the
 61    web UI's live table can react without polling.
 62    """
 63
 64    default_kind: ConsistKind = "software"
 65
 66    def __init__(
 67        self, throttles: ThrottleManager, system_prefix: str,
 68        programmers: ProgrammerManager | None = None,
 69    ) -> None:
 70        if not system_prefix:
 71            raise ValueError("system_prefix must be a non-empty string")
 72        self._throttles = throttles
 73        # Ops-mode programmer manager, used by the ``advanced`` kind to
 74        # write CV19 on each member decoder as they join or leave. May
 75        # be ``None`` on non-hardware setups (tests, InternalConnection);
 76        # advanced consists degrade silently to a no-op in that case
 77        # rather than raising.
 78        self._programmers: ProgrammerManager | None = programmers
 79        self.system_prefix = system_prefix
 80        self._consists: dict[int, Consist] = {}
 81        # Runtime drive state per consist address. Not persisted — the
 82        # consist itself is pure config (membership + wire kind), and
 83        # loco throttles carry the authoritative state. But the UI
 84        # needs a per-consist "last commanded" value to render the
 85        # slider position and highlight the active direction button;
 86        # without it every drive command re-renders the card at the
 87        # hardcoded defaults and the user sees the slider snap back to
 88        # 0 on every drag-release. Reset on delete.
 89        self._speeds: dict[int, float] = {}
 90        self._directions: dict[int, ThrottleDirection] = {}
 91        self.created: Signal[Consist] = Signal()
 92        self.deleted: Signal[Consist] = Signal()
 93        self.updated: Signal[Consist] = Signal()
 94
 95    def rebind_throttles(self, throttles: ThrottleManager) -> None:
 96        """Point the fan-out at a new :class:`ThrottleManager`.
 97
 98        Base :meth:`SystemConnection.__init__` builds the consist manager
 99        before drivers get a chance to install their hardware-aware
100        throttle manager in :meth:`open`. Without this hook the consist
101        manager keeps a stale reference to the pre-open placeholder and
102        every fan-out silently drives plain-Python :class:`Throttle`
103        beans that never touch the wire — the consist would look like
104        it's working (bean state advances, signals fire) while the
105        actual locomotive sits still. Drivers call this from ``open()``
106        immediately after replacing ``self.throttles``.
107        """
108        self._throttles = throttles
109
110    def rebind_programmers(
111        self, programmers: ProgrammerManager | None
112    ) -> None:
113        """Point the CV19 write path at a new :class:`ProgrammerManager`.
114
115        Same rationale as :meth:`rebind_throttles` — the driver's
116        hardware-aware programmer manager isn't in place until
117        :meth:`open` runs. Advanced consists need ops-mode access to
118        program CV19 on each member; without a real programmer the
119        writes silently become no-ops.
120        """
121        self._programmers = programmers
122
123    def current_speed(self, address: int) -> float:
124        """Last commanded fan-out speed for this consist, or 0.0."""
125        return self._speeds.get(address, 0.0)
126
127    def current_direction(self, address: int) -> ThrottleDirection:
128        """Last commanded direction for this consist, or FORWARD."""
129        return self._directions.get(address, ThrottleDirection.FORWARD)
130
131    # ---- CRUD ------------------------------------------------------------
132
133    async def create(
134        self,
135        address: int,
136        *,
137        kind: ConsistKind | None = None,
138        comment: str = "",
139    ) -> Consist:
140        """Create a new consist and register it. Idempotent by address.
141
142        Returns the existing consist if one already lives at *address*
143        rather than raising — matches JMRI's tolerance for repeated
144        create calls from restore paths.
145
146        Consist addresses are constrained to 1..127 short by the
147        ``Consist`` model; the caller doesn't need a ``long_address``
148        flag because there is no long-address form of a consist. See
149        :data:`jmri_core.beans.consist.CONSIST_ADDRESS_MAX`.
150        """
151        existing = self._consists.get(address)
152        if existing is not None:
153            return existing
154        consist = Consist(
155            address=address,
156            kind=kind or self.default_kind, comment=comment,
157        )
158        self._consists[address] = consist
159        await self._after_create(consist)
160        await self.created.emit(consist)
161        return consist
162
163    async def delete(self, address: int) -> bool:
164        """Remove the consist. Returns ``True`` if one was found."""
165        consist = self._consists.pop(address, None)
166        if consist is None:
167            return False
168        self._speeds.pop(address, None)
169        self._directions.pop(address, None)
170        await self._before_delete(consist)
171        await self.deleted.emit(consist)
172        return True
173
174    def get(self, address: int) -> Consist | None:
175        return self._consists.get(address)
176
177    def __iter__(self) -> Iterator[Consist]:
178        return iter(self._consists.values())
179
180    def __len__(self) -> int:
181        return len(self._consists)
182
183    def __contains__(self, address: int) -> bool:
184        return address in self._consists
185
186    # ---- membership -----------------------------------------------------
187
188    async def add_member(
189        self,
190        consist_address: int,
191        member_address: int,
192        *,
193        long_address: bool = False,
194        direction_normal: bool = True,
195    ) -> ConsistLoco:
196        """Add *member_address* to the consist at *consist_address*.
197
198        The member is appended at the end. If a loco with the same
199        address already exists in the consist, its direction flag is
200        updated in place and the existing :class:`ConsistLoco` is
201        returned. Fires :attr:`updated` on the consist so subscribers
202        can react to membership changes.
203        """
204        consist = self._consists.get(consist_address)
205        if consist is None:
206            raise KeyError(f"no consist at address {consist_address}")
207        for loco in consist.locos:
208            if loco.address == member_address:
209                loco.direction_normal = direction_normal
210                await self._after_membership_change(consist)
211                if consist.kind == "advanced":
212                    await self._program_cv19(
213                        member_address, long_address=loco.long_address,
214                        consist_address=consist_address,
215                        direction_normal=direction_normal,
216                    )
217                await self.updated.emit(consist)
218                return loco
219        loco = ConsistLoco(
220            address=member_address,
221            long_address=long_address,
222            direction_normal=direction_normal,
223            position=len(consist.locos),
224        )
225        consist.locos.append(loco)
226        await self._after_membership_change(consist)
227        if consist.kind == "advanced":
228            await self._program_cv19(
229                member_address, long_address=long_address,
230                consist_address=consist_address,
231                direction_normal=direction_normal,
232            )
233        await self.updated.emit(consist)
234        return loco
235
236    async def _program_cv19(
237        self,
238        member_address: int,
239        *,
240        long_address: bool,
241        consist_address: int,
242        direction_normal: bool,
243    ) -> None:
244        """Write CV19 on *member_address*'s decoder for an advanced
245        consist.
246
247        Per NMRA S-9.2.2, CV19 holds a 7-bit consist address in bits
248        0..6 and the direction flag in bit 7. Bit 7 set means the loco
249        runs in the reverse direction relative to the consist — i.e.
250        ``direction_normal=False``. A value of 0 in the low 7 bits
251        disables the consist for that decoder.
252        """
253        prog_mgr = self._programmers
254        if prog_mgr is None:
255            _log.debug(
256                "advanced consist %d: no programmer manager, skipping "
257                "CV19 write for member %d",
258                consist_address, member_address,
259            )
260            return
261        prog = prog_mgr.ops_programmer(
262            member_address, long_address=long_address,
263        )
264        if prog is None:
265            _log.debug(
266                "advanced consist %d: driver returned no ops "
267                "programmer for member %d; CV19 not written",
268                consist_address, member_address,
269            )
270            return
271        value = consist_address & 0x7F
272        if not direction_normal:
273            value |= 0x80
274        try:
275            await prog.write_cv(19, value)
276        except Exception as exc:  # noqa: BLE001
277            # POM writes can fail (no ack, decoder not on the track,
278            # RailCom garble). Log and keep going — the consist itself
279            # remains configured; the user can re-add the member after
280            # putting the loco on the rails.
281            _log.warning(
282                "CV19 write failed for consist %d, member %d "
283                "(value=%d): %s",
284                consist_address, member_address, value, exc,
285            )
286
287    async def _clear_cv19(
288        self, member_address: int, *, long_address: bool,
289        consist_address: int,
290    ) -> None:
291        """Reset CV19 to 0 on the ex-member's decoder so it stops
292        answering to the consist address."""
293        prog_mgr = self._programmers
294        if prog_mgr is None:
295            return
296        prog = prog_mgr.ops_programmer(
297            member_address, long_address=long_address,
298        )
299        if prog is None:
300            return
301        try:
302            await prog.write_cv(19, 0)
303        except Exception as exc:  # noqa: BLE001
304            _log.warning(
305                "CV19 clear failed for consist %d, ex-member %d: %s",
306                consist_address, member_address, exc,
307            )
308
309    async def remove_member(
310        self, consist_address: int, member_address: int,
311    ) -> bool:
312        """Drop *member_address* from the consist. Returns True if found.
313
314        Re-numbers surviving members' ``position`` fields so they stay
315        contiguous — the ordering is used for lead-loco / rear-loco
316        semantics in DCC-EX and NCE displays.
317        """
318        consist = self._consists.get(consist_address)
319        if consist is None:
320            return False
321        removed = False
322        removed_long_address = False
323        surviving: list[ConsistLoco] = []
324        for loco in consist.locos:
325            if loco.address == member_address and not removed:
326                removed = True
327                # Capture before the loco leaves — the ops programmer
328                # for the CV19 clear is keyed by the decoder's real
329                # address form (short vs long).
330                removed_long_address = loco.long_address
331                continue
332            surviving.append(loco)
333        if not removed:
334            return False
335        # Renumber positions.
336        for i, loco in enumerate(surviving):
337            loco.position = i
338        consist.locos = surviving
339        await self._after_membership_change(consist)
340        if consist.kind == "advanced":
341            await self._clear_cv19(
342                member_address, long_address=removed_long_address,
343                consist_address=consist_address,
344            )
345        await self.updated.emit(consist)
346        return True
347
348    # ---- fan-out throttle control ---------------------------------------
349
350    async def set_speed(self, consist_address: int, speed: float) -> None:
351        """Fan a speed command to every member throttle.
352
353        Members with ``direction_normal=False`` are treated identically
354        for speed (only direction is inverted). Missing throttles
355        (member never acquired) are silently acquired first so the
356        first speed command on a fresh consist works.
357
358        Records the commanded speed on the manager and emits ``updated``
359        so any subscriber (the /consists page's WebSocket bus in
360        particular) re-renders the card with the slider at the new
361        position instead of snapping back to zero.
362        """
363        consist = self._require(consist_address)
364        for member in consist.locos:
365            throttle = await self._throttles.acquire(
366                member.address, long_address=member.long_address,
367            )
368            await throttle.set_speed(speed)
369        self._speeds[consist_address] = speed
370        await self.updated.emit(consist)
371
372    async def set_direction(
373        self, consist_address: int, direction: ThrottleDirection,
374    ) -> None:
375        """Fan a direction command to every member, honouring per-loco
376        orientation.
377
378        A member with ``direction_normal=False`` is oriented backwards
379        in the consist, so it receives the *opposite* direction to keep
380        the whole train moving the same way — this is the core reason
381        consists exist at the software level.
382        """
383        consist = self._require(consist_address)
384        opposite = (
385            ThrottleDirection.REVERSE
386            if direction is ThrottleDirection.FORWARD
387            else ThrottleDirection.FORWARD
388        )
389        for member in consist.locos:
390            throttle = await self._throttles.acquire(
391                member.address, long_address=member.long_address,
392            )
393            wire = direction if member.direction_normal else opposite
394            await throttle.set_direction(wire)
395        self._directions[consist_address] = direction
396        await self.updated.emit(consist)
397
398    async def set_function(
399        self, consist_address: int, fn: int, on: bool,
400    ) -> None:
401        """Fan a function command to every member.
402
403        JMRI's convention is that all lit functions (F0/headlight, F1/
404        bell, F2/horn) fire on every loco; layout-specific automation
405        that wants function isolation should drive individual throttles
406        directly.
407        """
408        consist = self._require(consist_address)
409        for member in consist.locos:
410            throttle = await self._throttles.acquire(
411                member.address, long_address=member.long_address,
412            )
413            await throttle.set_function(fn, on)
414
415    async def emergency_stop(self, consist_address: int) -> None:
416        """Fan an emergency-stop to every member throttle."""
417        consist = self._require(consist_address)
418        for member in consist.locos:
419            throttle = await self._throttles.acquire(
420                member.address, long_address=member.long_address,
421            )
422            await throttle.emergency_stop()
423        # Reset UI speed so the slider snaps back to 0 after an e-stop
424        # rather than staying at whatever the user last commanded.
425        self._speeds[consist_address] = 0.0
426        await self.updated.emit(consist)
427
428    # ---- driver hooks ---------------------------------------------------
429
430    async def _after_create(self, consist: Consist) -> None:
431        """Called after a new consist is registered. Driver hook.
432
433        Subclasses install the wire-level group here — LocoNet's
434        ``LINK_SLOTS``, NCE's consist-memory write, DCC-EX's ``<C>``.
435        The base impl is a no-op (software consist).
436        """
437        return
438
439    async def _before_delete(self, consist: Consist) -> None:
440        """Called before a consist is dropped from the registry.
441
442        Subclasses tear the wire-level group down here.
443        """
444        return
445
446    async def _after_membership_change(self, consist: Consist) -> None:
447        """Called after ``add_member`` / ``remove_member`` mutates the
448        consist's locos list. Subclasses re-sync the wire-level group.
449        """
450        return
451
452    # ---- internals ------------------------------------------------------
453
454    def _require(self, address: int) -> Consist:
455        consist = self._consists.get(address)
456        if consist is None:
457            raise KeyError(f"no consist at address {address}")
458        return consist