Vmux signal types

src/fixate/_switching.py

@@ -43,17 +43,23 @@
     Dict,
     FrozenSet,
     Iterable,
+    Literal,
+    get_origin,
+    get_args,
 )
 from dataclasses import dataclass
 from functools import reduce
 from operator import or_
+from types import new_class
 
 Signal = str
+EmptySignal = Literal[""]
 Pin = str
 PinList = Sequence[Pin]
 PinSet = FrozenSet[Pin]
+# do we bother to add EmptySignal here?
 SignalMap = Dict[Signal, PinSet]
-TreeDef = Sequence[Union[Signal, "TreeDef"]]
+TreeDef = Sequence[Union[Optional[Signal], "TreeDef"]]
 
 
 @dataclass(frozen=True)
@@ -87,14 +93,73 @@ def __or__(self, other: PinUpdate) -> PinUpdate:
 
 PinUpdateCallback = Callable[[PinUpdate, bool], None]
 
+S = TypeVar("S", bound=str)
 
-class VirtualMux:
+
+class VirtualMux(Generic[S]):
+    map_tree: Optional[TreeDef] = None
+    map_list: Optional[Sequence[Sequence[str]]] = None
     pin_list: PinList = ()
     clearing_time: float = 0.0
 
+    def __class_getitem__(cls, arg):
+        # https://peps.python.org/pep-0560
+
+        # so the problem we have to solve is the way this works is convoluted and the __orig_bases__ property we rely on is not preserved
+        # __orig_bases__ is an attribute set by the metaclass that is used to store the type information provided by __class_getitem__
+        # expected class creation:
+        # VirtualMux[type] -> VirtualMux.__class_getitem__(type) -> VirtualMux.__init__() (with type info visible on __orig_bases__)
+        # what actually happens
+        # VirtualMux[type] -> VirtualMux.__class_getitem__(type) -> _GenericAlias(VirtualMux).__call__() -> VirtualMux.__init__()
+        # the _GenericAlias middle class means we lose the type information provided through the __class_getitem__ call
+
+        # the attribute __orig_class__ can be used to store the original class that was created with __class_getitem__
+        # HOWEVER this relies on the __init__ of VirtualMux succeeding - we end up in a circular dependency of needing
+        # the init to succeed to set the attribute, but needing the attribute to be set to succeed in the init (due to our implementation)
+
+        # there are two options
+        # 1. create a proxy class and wrap it to look like a normal VirtualMux
+        # 2. dynamically create the pin_list map_list here
+
+        proxy = new_class(f"{cls}[{arg}]", bases=(cls,))
+        # two cases: either we are a passing in the generic typevar, or we are passing in the mux definition
+        if type(arg) != TypeVar:
+            pin_list, map_list = cls._unpack_muxdef(arg)
+            proxy.pin_list = pin_list
+            proxy.map_list = map_list
+        return proxy  # now the actual class can be initialised
+
+    @staticmethod
+    def _unpack_muxdef(muxdef):
+        # two cases
+        # muxdef is the signal definition
+        assert get_origin(muxdef) == Union, "MuxDef must be a Union of signals"
+        signals = get_args(muxdef)
+        map_list: Sequence[Sequence[str]] = []
+        pin_list: PinList = []
+        for s in signals:
+            # get_args gives Literal
+            # get_args ignores metadata before 3.10
+            sigdef, *pins = get_args(s)
+            # 3.8 only
+            if not pins:
+                if not hasattr(s, "__metadata__"):
+                    raise ValueError(
+                        "VirtualMux Subclass must define the pins for each signal"
+                    )
+                # s.__metadata__ is our pin list
+                pins = s.__metadata__
+            # get_args gives members of Literal
+            (signame,) = get_args(sigdef)
+            assert isinstance(signame, Signal), "Signal name must be signal type"
+            assert all(isinstance(p, Pin) for p in pins), "Pins must be pin type"
+            pin_list.extend(pins)
+            map_list.append((signame, *pins))
+
+        return pin_list, map_list
+
     ###########################################################################
     # These methods are the public API for the class
-
     def __init__(self, update_pins: Optional[PinUpdateCallback] = None):
         self._last_update_time = time.monotonic()
 
@@ -104,16 +169,16 @@ def __init__(self, update_pins: Optional[PinUpdateCallback] = None):
         else:
             self._update_pins = update_pins
 
+        self._pin_set = frozenset(self.pin_list)
+        self._signal_map: SignalMap = self._map_signals()
+
         # We annotate the pin_list to be an ordered sequence, because if the
         # mux is defined with a map_tree, we need the ordering. But after
         # initialisation, we only need set operations on the pin list, so
         # we convert here and keep a reference to the set for future use.
-        self._pin_set = frozenset(self.pin_list)
 
         self._state = ""
 
-        self._signal_map: SignalMap = self._map_signals()
-
         # Define the implicit signal "" which can be used to turn off all pins.
         # If the signal map already has this defined, raise an error. In the old
         # implementation, it allows the map to set this, but when switching the
@@ -129,7 +194,9 @@ def __init__(self, update_pins: Optional[PinUpdateCallback] = None):
         if hasattr(self, "default_signal"):
             raise ValueError("'default_signal' should not be set on a VirtualMux")
 
-    def __call__(self, signal: Signal, trigger_update: bool = True) -> None:
+    def __call__(
+        self, signal: Union[S, EmptySignal], trigger_update: bool = True
+    ) -> None:
         """
         Convenience to avoid having to type jig.mux.<MuxName>.multiplex.
 
@@ -138,7 +205,9 @@ def __call__(self, signal: Signal, trigger_update: bool = True) -> None:
         """
         self.multiplex(signal, trigger_update)
 
-    def multiplex(self, signal: Signal, trigger_update: bool = True) -> None:
+    def multiplex(
+        self, signal: Union[S, EmptySignal], trigger_update: bool = True
+    ) -> None:
         """
         Update the multiplexer state to signal.
 
@@ -230,13 +299,13 @@ def _map_signals(self) -> SignalMap:
         Avoid subclassing. Consider creating helper functions to build
         map_tree or map_list.
         """
-        if hasattr(self, "map_tree"):
+        if self.map_tree is not None:
             return self._map_tree(self.map_tree, self.pin_list, fixed_pins=frozenset())
-        elif hasattr(self, "map_list"):
+        elif self.map_list is not None:
             return {sig: frozenset(pins) for sig, *pins in self.map_list}
         else:
             raise ValueError(
-                "VirtualMux subclass must define either map_tree or map_list"
+                "VirtualMux subclass must define either map_tree or map_list or provide a type to VirtualMux"
             )
 
     def _map_tree(self, tree: TreeDef, pins: PinList, fixed_pins: PinSet) -> SignalMap:
@@ -482,7 +551,7 @@ def __init__(
         super().__init__(update_pins)
 
 
-class RelayMatrixMux(VirtualMux):
+class RelayMatrixMux(VirtualMux[S]):
     clearing_time = 0.01
 
     def _calculate_pins(

src/fixate/examples/jig_driver.py

@@ -10,6 +10,7 @@
     MuxGroup,
     PinValueAddressHandler,
     VirtualSwitch,
+    RelayMatrixMux,
 )
 
 
@@ -58,3 +59,70 @@ class JigMuxGroup(MuxGroup):
 jig.mux.mux_two("sig5")
 jig.mux.mux_three("On")
 jig.mux.mux_three(False)
+
+
+from typing_extensions import Annotated
+
+from typing import Literal, Union
+
+# maybe we can create aliases to make it easier to understand how to create a MuxDef
+# this doesn't work well because these are instances of type, not instances of object
+SignalName = Literal
+Signal = Annotated
+MuxDef = Union
+
+# fmt: off
+MuxOneSigDef = MuxDef[
+    Signal[SignalName["sig_a1"], "a0", "a2"], 
+    Signal[SignalName["sig_a2"], "a1"]
+]
+# fmt: on
+rmm = RelayMatrixMux[MuxOneSigDef]()
+
+from typing import TypeVar
+
+S = TypeVar("S", bound=str)
+
+
+class MuxA(VirtualMux[MuxOneSigDef]):
+    """A mux definition used by a few scripts"""
+
+
+muxa = MuxA()
+muxa("sig_a2")
+muxa("sig_a1")
+
+
+class MuxC(VirtualMux[S]):
+    ...
+
+
+muxc = MuxC[MuxOneSigDef]()
+muxc("sig_a2")
+muxc("sig_a1")
+muxa("sig_a")
+
+muxb = VirtualMux[MuxOneSigDef]()
+
+muxb.multiplex("sig_a2")
+muxb.multiplex("sig_a1")
+try:
+    muxb.multiplex("1")
+except ValueError:
+    ...
+else:
+    raise ValueError("muxb.multiplex('1') should have raised a ValueError")
+
+# an example of generic subclasses of VirtualMux
+rmm = RelayMatrixMux[MuxOneSigDef]()
+
+rmm("sig_a1")
+
+try:
+    rmm("sig")
+except ValueError:
+    ...
+else:
+    raise ValueError("rmm('sig') should have raised a ValueError")
+
+rmm("fuck")

test/test_switching.py

@@ -15,6 +15,9 @@
     JigDriver,
 )
 
+from typing import Literal, Union, TypeVar, Generic
+from typing_extensions import Annotated  # only for 3.8
+
 import pytest
 
 ################################################################
@@ -592,3 +595,56 @@ def test_pin_update_or():
         2.0,
     )
     assert expected == a | b
+
+
+# fmt: off
+MuxASigDef = Union[
+    Annotated[Literal["sig_a1"], "a0", "a1"], 
+    Annotated[Literal["sig_a2"], "a1"]
+]
+# fmt: on
+
+
+def test_typed_mux():
+    clear = PinSetState(off=frozenset({"a0", "a1"}))
+    a1 = PinSetState(on=frozenset({"a0", "a1"}))
+    a2 = PinSetState(on=frozenset({"a1"}), off=frozenset({"a0"}))
+
+    updates = []
+    updatesa = []
+
+    mux = VirtualMux[MuxASigDef](lambda x, y: updates.append((x, y)))
+    mux_a = MuxA(lambda x, y: updatesa.append((x, y)))
+    assert mux_a._signal_map == mux._signal_map
+    assert mux_a._pin_set == mux._pin_set
+
+    mux("sig_a1")
+    mux_a("sig_a1")
+    assert updates.pop() == updatesa.pop() == (PinUpdate(PinSetState(), a1), True)
+
+    mux.multiplex("sig_a2", trigger_update=False)
+    mux_a.multiplex("sig_a2", trigger_update=False)
+    assert updates.pop() == updatesa.pop() == (PinUpdate(PinSetState(), a2), False)
+
+    mux("")
+    mux_a("")
+    assert updates.pop() == updatesa.pop() == (PinUpdate(PinSetState(), clear), True)
+
+
+def test_typed_mux_subclass():
+    class SubMux(VirtualMux[MuxASigDef]):
+        pass
+
+    sm = SubMux()
+    assert sm._signal_map == MuxA()._signal_map
+    assert sm._pin_set == MuxA()._pin_set
+
+
+def test_typed_mux_generic_subclass():
+    T = TypeVar("T", bound=str)
+
+    class GenericSubMux(VirtualMux[T]):
+        pass
+
+    gsm = GenericSubMux[MuxASigDef]()
+    assert gsm._signal_map == MuxA()._signal_map