Alternate concept for VirtualMux

[clint-lawrence]

This came about while thinking of other ways to create a VirtualMux that will be type hint and IDE friendly. Originally a comment on #184 but I think it is better to have an issue separate from that so it doesn't get lost.

Each signal would be a method on the resulting VirtualMux subclass. This would be a big change for existing scripts. But once the jig refactor PR lands, we could introduce something like this and scripts could choose which style to use. For new scripts, the IDE autocomplete benefits of this approach would be nice to have. (There might be a way to type-hint the current VirtualMux, but I haven't figured it out)

# existing
class Mux(VirtualMux):
    pin_list = ("x0", "x1", "x2")
    map_list = (
        ("sig1", "x0"),
        ("sig2", "x1"),
        ("sig3", "x2", "x0"),
    )

#new
@dataclass
class MuxNew(VirtualMux):
    sig1: Sig = signal("x0")
    sig2: Sig = signal("x1")
    sig3: Sig = signal("x2", "x0")

dm.jig.mux.Mux("sig2")
# becomes
dm.jig.mux.NewMux.sig2()

# For tree defined multiplexers
#                                   __________
# a0-------------------------------|          |
#                       ________   |          |
# a1_b0----------------|        |--|          |
# a1_b1----------------| Mux B  |  |          |
#              (None)--|  4:1   |  |          |
# a1_b3----------------| x3  x2 |  |          |
#                      |________|  |          |
#                                  |  Mux A   |
# a2-------------------------------|   4:1    |
#                       ________   |          |
# a3_c0----------------|        |--|          |
# a3_c1----------------| Mux C  |  |          |
# a3_c2----------------|  4:1   |  |          |
#              (None)--| x5  x4 |  |          |
#                      |________|  |          |
#                                  |__x1__x0__|
muxA = MuxTree("x0", "x1")
muxB = MuxTree("x2", "x3", nest=muxA)
muxC = MuxTree("x4", "x5", nest=muxA)

@dataclass
class Mux(VirtualMux):
    a0: Sig = muxA()
    a1_b0: Sig = muxB()
    a1_b1: Sig = muxB()
    muxB()  # "throw away" one entry
    a1_b3: Sig = muxB()
    a2: Sig = muxA()
    a3_c0: Sig = muxC()
    a3_c1: Sig = muxC()
    a3_c2: Sig = muxC()

Implementation - Maybe? very not tested... Most likely bugs, but I think the concept is O.K. I think we would end up with signal returning a class instead of a function, but this maps out the general idea.

Sig = Callable[[VirtualMux, bool], None]

def signal(*pins) -> Sig:
    def _func(self: VirtualMux, trigger_update: bool=False):
        self._do_switch(pins, trigger_update)
    _func.pins = pins
    return _func


class MuxTree:
    def __init__(self, *pins: tuple[Pin], nest: MuxTree) -> None:
        self._pins = set(pins)
        self._pin_gen = generate_bit_sets(pins)
        if nest is not None:
            self._fixed_pins = nest().pins
        else:
            self._fixed_pins = set()

    def __call__(self) -> Sig:
        pins = self._fixed_pins  | next(self._pin_gen)
        return signal(*pins)

[clint-lawrence]

@daniel-montanari?

S = TypeVar("S")
class VirtualMux(Generic[S]):
    def __init__(self):
        self._signal_map: dict[Signal, set[Pin]] = {}

    def __call__(self, signal: S, trigger_update: bool=False) -> None:
        self.multiplex(signal, trigger_update)
    
    def multiplex(self, signal: S, trigger_update: bool=False) -> None:
        print(self._signal_map[signal])

MuxOneSigDef = Union[
    Annotated[Literal["sig1"], ("x0",)],
    Annotated[Literal["sig2"], ("x1",)],
    Annotated[Literal["sig3"], ("x0", "x1")],
]

class MuxOne(VirtualMux[MuxOneSigDef]):
    pass

[clint-lawrence]

Possible way to extend this idea to "Tree" mux definitions.

Unravelling this in code to build the signal map would be a bit of a mess, but I think it captures all the required info.

There is one weird thing. I'm stealing the "" literal as a place holder for "check the metadata". It is always a valid signal, but defined to mean "all pins off". So we can use it here as a place holder for either a mux input that isn't used, or as a mux input that "connects" to another mux. Distinguishing between those two is a matter of context and the second metadata value (i.e. the tuple of (MuxB, MuxC) in the definition of MuxA below).

For each instance of Annotated there are two or three values:

• A list of literal strings, corresponding to signals
• A tuple of pins
• optionally a tuple of other type definitions or None

MuxC = Annotated[Literal["c0", "c1", "c2"], ("x4", "x5")]
MuxB = Annotated[Literal["b0", "b1", "", "b3"], ("x2", "x3"), (None,)]  # None corresponds to an unused signal, marked by "" 
MuxA = Annotated[Literal["a0", "", "a2", ""], ("x0", "x1"), (MuxB, MuxC)]  # MuxB & MuxC connect to the entries marked by ""
MuxTreeDef = Union[MuxA, MuxB, MuxC]

class TreeDefinedMux[MuxTreeDef]:
    pass

[clint-lawrence]

Need to check if this has an impact on that tree definition idea, because it is relying on ordering
https://docs.python.org/3/library/typing.html#typing.get_args

If X is a union or Literal contained in another generic type, the order of (Y, Z, ...) may be different from the order of the original arguments [Y, Z, ...] due to type caching. Return () for unsupported objects.

[daniel-montanari]

    def _map_signals(self) -> SignalMap:
        """
        Default implementation of the signal mapping

        We need to construct a dictionary mapping signals to a set of pins.
        In the case that self.map_list is set, the is pretty trival.
        If the mux is defined with self.map_tree we have more work to
        do, which is recursively delegated to _map_tree

        Avoid subclassing. Consider creating helper functions to build
        map_tree or map_list.
        """
        if hasattr(self, "map_tree"):
            return self._map_tree(self.map_tree, self.pin_list, fixed_pins=frozenset())
        elif hasattr(self, "map_list"):
            return {sig: frozenset(pins) for sig, *pins in self.map_list}
        elif (self.__orig_bases__ != VirtualMux.__orig_bases__):
            # the user has provided map_list using annotations
            # if the type annotations have not been supplied, then self.__orig_bases__ == VirtualMux.__orig_bases__
            # if they have been overridden, then self.__orig_bases__ != VirtualMux.__orig_bases__
            # works in 3.8 and 3.12
            # this will only work if VirtualMux was subclassed
            # if creating an instance using VirtualMux[type]() then __orig_bases__ does not exist
            # there seems to be some

            # we are expecting exactly 1 value, unpack it
            cls, = self.__orig_bases__
            assert get_origin(cls) == VirtualMux
            lst, = get_args(cls)
            assert get_origin(lst) == Union
            signals = get_args(lst)
            sigmap = {}
            for s in signals:
                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 Literal
                literal, = get_args(s)
                # get_args gives members of Literal
                signame, = get_args(literal)
                sigmap[signame] = frozenset(pins)

            return sigmap
        else:
            raise ValueError(
                "VirtualMux subclass must define either map_tree or map_list or provide a type to VirtualMux"
            )

This nearly gets there. It'd be nice if we could do something like

MuxOne = VirtualMux[MuxOneSigDef]

@dataclass
class JigMuxGroup(MuxGroup):
    mux_one: MuxOne = field(default_factory=MuxOne)

I think that would get rid of my weird feelings about the way we subclass a thing to do basically nothing. But I get that this whole idea is bastardising the crap out of this system anyway.

The __orig_bases__ trick only works due to subclassing. You'd need another way to bind the argument of __class_getitem__ to the associated __init__

[daniel-montanari]

    def __class_getitem__(cls, *args, **kwargs):
        # without calling getitem the class doesn't work as a generic
        getitm =  super().__class_getitem__(*args, **kwargs) # normally returns a generic
        # create a proxy to force the __init_subclass__ hook
        class Hack(getitm):
            ...

        return Hack # now the actual class can be initialised

Forgive me for I have sinned.
Type checkers ignore this completely and just use the inferred behaviour due to subclassing Generic. We could use TYPE_CHECKING if we want to lie to the type checker as well.

[daniel-montanari]

Looks like they changed how Annotated behaves with get_args between 3.8 and 3.10.
In 3.8 __metadata__ is ignored by get_args. In 3.10 it is used. I guess this makes sense if get_args is meant to return whatever is inside the brackets of a generic.
Changing the unpack operator to allow variable length is enough to get this working on all versions.

[daniel-montanari]

Made a PR for people to look at and play around with.
Not expecting it to be merged in any time soon.