enum.py

def enum(**variants):
    def enhance(cls):
        def make_constructor(tag, sig):
            def constructor(*data):
                nonlocal sig
                if callable(sig):
                    sig = sig()

                if len(sig) != len(data):
                    raise ValueError(f"expected {len(sig)} items, not {len(data)}")
                for x, ty in zip(data, sig):
                    if not isinstance(x, ty):
                        raise TypeError(f"expected {ty} but got {repr(x)}")
                inst = cls()
                inst.tag = tag
                inst.data = data
                return inst

            return constructor

        for tag, sig in variants.items():
            setattr(cls, tag, make_constructor(tag, sig))

        def match(self, **handlers):
            if self.tag in handlers:
                return handlers[self.tag](*self.data)
            elif "_" in handlers:
                return handlers["_"]()
            else:
                raise ValueError(f"missing handler for {self.tag}")

        setattr(cls, "match", match)

        def _repr(self):
            return f"{cls.__name__}.{self.tag}({', '.join(map(repr, self.data))})"

        setattr(cls, "__repr__", _repr)

        def _eq(self, other):
            if not isinstance(other, cls):
                return NotImplemented
            return self.tag == other.tag and self.data == other.data

        setattr(cls, "__eq__", _eq)

        return cls

    return enhance