Vmux signal types

mypy.ini

@@ -90,4 +90,10 @@ follow_imports = silent
 [mypy-fixate.drivers,fixate.drivers.dso.helper,fixate.drivers.funcgen.helper,fixate.drivers.funcgen.rigol_dg1022,fixate.drivers.pps,fixate.drivers.pps.helper,fixate.drivers.ftdi]
 follow_imports = silent
 [mypy-fixate.examples.function_generator,fixate.examples.programmable_power_supply,fixate.examples.test_script]
-follow_imports = silent
+follow_imports = silent
+
+# this example demos the type checker warning you when using the wrong arguments
+# we explicitly suppress the errors in the script, but make sure that mypy
+# actually found an error
+[mypy-fixate.examples.jig_driver]
+enable_error_code = unused-ignore

src/fixate/_switching.py

@@ -43,17 +43,33 @@
     Dict,
     FrozenSet,
     Iterable,
+    Literal,
+    get_origin,
+    get_args,
+    Any,
 )
 from dataclasses import dataclass
 from functools import reduce
 from operator import or_
+from types import new_class
+
+import sys
+
+if sys.version_info >= (3, 9):
+    from typing import Annotated
+else:
+    from typing_extensions import Annotated
+
 
 Signal = str
+EmptySignal = Literal[""]
 Pin = str
 PinList = Sequence[Pin]
 PinSet = FrozenSet[Pin]
+MapList = Sequence[Sequence[Union[Signal, 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,15 +103,23 @@ def __or__(self, other: PinUpdate) -> PinUpdate:
 
 PinUpdateCallback = Callable[[PinUpdate, bool], None]
 
+S = TypeVar("S", bound=str)
 
-class VirtualMux:
+from types import get_original_bases, resolve_bases
+
+
+class VirtualMux(Generic[S]):
+    map_tree: Optional[TreeDef] = None
+    map_list: Optional[Sequence[Sequence[str]]] = None
     pin_list: PinList = ()
     clearing_time: float = 0.0
 
     ###########################################################################
     # These methods are the public API for the class
-
     def __init__(self, update_pins: Optional[PinUpdateCallback] = None):
+        # digest all the typing information if there is any to sete pin_list and map_list
+        self._digest_type_hints()
+
         self._last_update_time = time.monotonic()
 
         self._update_pins: PinUpdateCallback
@@ -104,16 +128,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 +153,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 +164,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 +258,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:
@@ -441,6 +469,50 @@ def _default_update_pins(
         """
         print(pin_updates, trigger_update)
 
+    def _digest_type_hints(self) -> None:
+        # digest all the typing information if there is any
+        bases = get_original_bases(self.__class__)
+        resolved_bases = resolve_bases(bases)
+        first_resolved_base = resolved_bases[0]
+        assert issubclass(
+            first_resolved_base, VirtualMux
+        ), f"{first_resolved_base} should be VirtualMux subclass"
+        if bases != resolved_bases:
+            args = get_args(bases[0])
+            plist, mlist = self._unpack_muxdef(args[0])
+            self.pin_list = plist
+            self.map_list = mlist
+
+    @staticmethod
+    def _unpack_muxdef(muxdef: type) -> tuple[PinList, MapList]:
+        # muxdef is the signal definition
+        if get_origin(muxdef) == Union:
+            signals = get_args(muxdef)
+        elif get_origin(muxdef) == Annotated:
+            # Union FORCES you to have two or more types, so this handles the case of only one pin
+            signals = (muxdef,)
+        else:
+            raise TypeError("Signal definition must be Union or Annotated")
+
+        map_list: list[tuple[str]] = []
+        pin_list: list[str] = []
+        for s in signals:
+            # not working, get_origin(Annotated) returns the wrapped type in 3.8
+            # assert get_origin(s) == Annotated, "Signal definition must be annotated"
+            # get_args gives Literal
+            sigdef, *pins = get_args(s)
+            assert (
+                get_origin(sigdef) == Literal
+            ), "Signal definition must be string literal"
+            # 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
+
 
 class VirtualSwitch(VirtualMux):
     """
@@ -482,7 +554,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,87 @@ class JigMuxGroup(MuxGroup):
 jig.mux.mux_two("sig5")
 jig.mux.mux_three("On")
 jig.mux.mux_three(False)
+
+
+# VirtualMuxes can be created with type annotations to provide the signal map
+from typing import Literal, Annotated, Union
+
+# a signal is a typing Annotation
+# the first Literal is the signal name, the rest are the pin names
+# the signal name MUST be a Literal
+# multiple signals can be combined with a Union
+# assigning annotations to variables is possible
+# fmt: off
+MuxOneSigDef = Union[
+    Annotated[Literal["sig_a1"], "a0", "a2"], 
+    Annotated[Literal["sig_a2"], "a1"],
+]
+
+MuxTwoSigDef = Union[
+    Annotated[Literal["sig_b1"], "b0", "b2"],
+    Annotated[Literal["sig_b2"], "b1"],
+]
+
+# if defining only a single signal, the Union is omitted in the definition
+SingleSingleDef = Annotated[Literal["sig_c1"], "c0", "c1"]
+# fmt: on
+
+# VirtualMuxes can now be created with type annotations to provide the signal map
+# this only works when subclassing
+class MyMux(VirtualMux[MuxOneSigDef]):
+    """A helpful description for my mux that is used in this jig driver"""
+
+
+muxa = MyMux()
+
+muxa("sig_a1")
+muxa("sig_a2")
+
+# using the wrong signal name will be caught at runtime and by the type checker
+try:
+    muxa("unknown signal name")  # type: ignore[arg-type]
+except ValueError as e:
+    print(f"An Exception would have occurred: {e}")
+else:
+    raise ValueError("Should have raised an exception")
+
+
+class MultiPinSwitch(VirtualMux[SingleSingleDef]):
+    """This acts like a switch, but has to coordinate two pins"""
+
+
+ls = MultiPinSwitch()
+ls("sig_c1")
+ls("")
+
+# further generic types can be created by subclassing from VirtualMux using a TypeVar
+# compared to the above way of subclassing, this way lets you reuse the class
+
+from typing import TypeVar
+
+S = TypeVar("S", bound=str)
+
+
+class MyGenericMux(VirtualMux[S]):
+    ...
+
+    def extra_method(self) -> None:
+        print("Extra method")
+
+
+class MyConcreteMux(MyGenericMux[MuxTwoSigDef]):
+    pass
+
+
+generic_mux = MyConcreteMux()
+generic_mux("sig_b1")
+generic_mux("sig_b2")
+
+# RelayMatrixMux is an example of this
+class MyRelayMatrixMux(RelayMatrixMux[MuxOneSigDef]):
+    pass
+
+
+rmm = MyRelayMatrixMux()
+rmm("sig_a1")
+rmm("sig_a2")