match.py

#!/usr/bin/env python3

import enum


def match(v, allow_partial=False, allow_impossible=False):
    """Provides exhaustive pattern matching for Enum."""
    if not isinstance(v, enum.Enum):
        raise ValueError(f'Supplied value is not an Enum: {v}')

    def go(**kwargs):
        if not allow_partial or not allow_impossible:
            all_cases = set(x.name for x in v.__class__)
            supplied_cases = set(kwargs)
            if not allow_impossible:
                impossible = supplied_cases.difference(all_cases)
                if impossible:
                    raise TypeError(f'Impossible cases supplied: {impossible}')
            if not allow_partial:
                missing = all_cases.difference(supplied_cases)
                if missing:
                    raise TypeError(f'Missing cases: {missing}')
        f = kwargs.get(v.name)
        if not v:
            raise MatchError(v)
        return f()

    return go

class MatchError(Exception):
    """Exception thrown due to a 'match' failure."""
    pass

class Suit(enum.Enum):
    """Example enum to demonstrate 'match'."""
    Clubs = enum.auto()
    Hearts = enum.auto()
    Spades = enum.auto()
    Diamonds = enum.auto()


def main():
    # Happy path. We create a render function that uses 'match' in such
    # a way that we must supply exactly all the cases for the enum.

    print('render (happy path)')
    print('-----------')

    def render(suit):
        return match(suit)(
            Clubs=lambda: '♣',
            Hearts=lambda: '♥',
            Spades=lambda: '♠',
            Diamonds=lambda: '♦',
        )

    for s in Suit:
        print(f'render({s}) = ', end='')
        print(render(s))

    # Sad path. We create a render that doesn't handle the 'Diamonds'
    # case. It will fail no matter what value we pass it.

    print('')
    print('render (missing cases)')
    print('----------')

    def render(suit):
        return match(suit)(
            Clubs=lambda: '♣',
            Hearts=lambda: '♥',
            Spades=lambda: '♠',
        )

    for s in Suit:
        try:
            print(f'render({s}) = ', end='')
            print(render(s))
        except Exception as e:
            print(repr(e))

    # We create a render that doesn't handle the 'Diamonds' case, but we give
    # it 'allow_partial=True' to say this is ok. It only fails when supplied
    # the unhandled case(s).

    print('')
    print('render (allow_partial=True)')
    print('----------')

    def render(suit):
        return match(suit, allow_partial=True)(
            Clubs=lambda: '♣',
            Hearts=lambda: '♥',
            Spades=lambda: '♠',
        )

    for s in Suit:
        try:
            print(f'render({s}) = ', end='')
            print(render(s))
        except Exception as e:
            print(repr(e))

    # Sad path. We create a render that handles all the cases but also handles
    # a case that that can never happen, 'Star'. By default, this will error
    # for any supplied value.

    print('')
    print('render (impossible cases)')
    print('----------')

    def render(suit):
        return match(suit)(
            Stars=lambda: '★',
            Clubs=lambda: '♣',
            Hearts=lambda: '♥',
            Spades=lambda: '♠',
            Diamonds=lambda: '♦',
        )

    for s in Suit:
        try:
            print(f'render({s}) = ', end='')
            print(render(s))
        except Exception as e:
            print(repr(e))

    # We create a render that handles all the cases but also handles
    # a case that that can never happen, 'Star'. We supply
    # 'allow_impossible=True' to allow this.

    print('')
    print('render (allow_impossible=True)')
    print('----------')

    def render(suit):
        return match(suit, allow_impossible=True)(
            Stars=lambda: '★',
            Clubs=lambda: '♣',
            Hearts=lambda: '♥',
            Spades=lambda: '♠',
            Diamonds=lambda: '♦',
        )

    for s in Suit:
        try:
            print(f'render({s}) = ', end='')
            print(render(s))
        except Exception as e:
            print(repr(e))


if __name__ == '__main__':
    main()