how to use Python multimethod with custom class as argument type

212 Views Asked by At

I have the following classes setup:

class A:
    pass

class B(A):
    pass

class C(A):
    pass

Then I have a method that gets called with 2 arguments:

  • x -> always a string
  • y -> one of the classes above (A, B, or C) And we do something different depending on which class was passed as second argument.
@multimethod
def click(x: str, y: A=None):
    do_A_thing()

@multimethod
def click(x: str, y: B=None):
    do_B_thing()

@multimethod
def click(x: str, y: object=None):
    raise expeption

With this configuration, an exception is always raised even when y is of type A or B.

I also looked at using singledispatch, but I need the validation on the second argument and I can't switch their places. I also tried using multipledispatch, but it didn't work similarly to multimethod.

1

There are 1 best solutions below

3
On

Have a look to the doc. Notice further that the dispatch occurs "on the type of the first argument". So you should refactor the signature of the functions.

Here a minimal working example for a functions annotated with types

@functools.singledispatch
def click(y:A, x:str):
    print('A')

@click.register
def _(y:B, x:str, ):
    print('B')

@click.register
def _(y:C, x:str):
    print('C')

click(B(), "x")
#B

and for a function which doesn’t use type annotations you need to pass the appropriate type argument explicitly to the decorator itself

@functools.singledispatch
def click(y, x):
    print('A')

@click.register(B)
def _(y, x):
    print('B')

@click.register(C)
def _(y, x):
    print('C')

x = 'str'
click(B(), x)
#B

For switching the order of the argument it's enough to define a decorator, arg_swapper, with this functionality

def arg_swapper(func):
    # interchange the order of the args of a function (assuming it takes exactly two parameters!)
    def wrapper(func_arg1, func_arg2):
        return func(func_arg2, func_arg1)
    return wrapper


def do_A_thing(x, y): print(f"do_A_thing({x=},{y=})")
def do_B_thing(x, y): print(f"do_B_thing({x=},{y=})")
def do_C_thing(x, y): print(f"do_C_thing({x=},{y=})")


@functools.singledispatch
@arg_swapper
def click(x, y): do_A_thing(x, y)    

@click.register(B)
@arg_swapper
def _(x, y): do_B_thing(x, y)

@click.register(C)
@arg_swapper
def _(x, y): do_C_thing(x, y)

x = 'str'
click(B(), x)
click(A(), x)
click(C(), x)

# Output
#do_B_thing(x='string',y=<__main__.B object at 0x7f8ef577bd10>)
#do_A_thing(x='string',y=<__main__.A object at 0x7f8ef577bd10>)
#do_C_thing(x='string',y=<__main__.C object at 0x7f8ef577bd10>)

Notice that you either stack the decorators (as shown above) or do an explicit call

@click.register(C)
def _(x, y): arg_swapper(do_C_thing)(x, y)