Function Overload In Python With Single Dispatch Python

Sep 17th, 2021 - written by Kimserey with .

Python supports overloading functions via single dispatch. With single dispatch we can define multiple implementations of a function and have it chosen based on a single argument. In this post we will look at how single dispatch works with example.

singledispatch

singledispatch is provided by functools:

1
>>> from functools import singledispatch

We can then register a function for overload with @singledispatch decorator:

1
2
3
4
In [11]: @singledispatch
    ...: ... def my_func(arg1, arg2):
    ...:         print("default", arg1, arg2)
    ...: 

Then with the name of the function we can register overloads:

1
2
3
4
5
6
7
8
9
10
11
12
13
In [24]: @my_func.register
    ...: def _(arg1: str, arg2):
    ...:     print("arg1 str", arg1, arg2)

In [25]: @my_func.register
    ...: def _(arg1: int, arg2):
    ...:     print("arg1 int", arg1, arg2)

In [26]: my_func(1, "hello")
arg1 int 1 hello

In [27]: my_func("test", "hello")
arg1 str test hello

Here we can see that the type annotation :int and :str was used to choose the implementation based on the first argument type. If a code isn’t using type annotation, we can directly provide a type to register:

1
2
3
4
5
6
In [31]: @my_func.register(float)
    ...: def _(arg1, arg2):
    ...:     print("arg float", arg1, arg2)

In [33]: my_func(1.0, "hello")
arg float 1.0 hello

Lastly we can specify multiple overload with the same implementation by registering multiple times:

1
2
3
4
5
6
7
8
9
10
In [46]: @my_func.register(float)
    ...: @my_func.register(list)
    ...: def _(arg1, arg2):
    ...:     print("arg float list", arg1, arg2)

In [47]: my_func(1.0, 1)
arg float list 1.0 1

In [48]: my_func([1,2,3], 1)
arg float list [1, 2, 3] 

Note that the dispatch only looks at the type of the first argument, and register only takes one cls argument which refers to the first argument.

singledispatchmethod

In the same way, we can overload method classes with singledispatchmethod. The difference between singledispatch and singledispatchmethod is that the dispatch happens on the type of the first non-self or non-cls argument. This is so that both method and class-method can be overloaded:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
In [4]: class Hi:
   ...:     @singledispatchmethod
   ...:     def say(self, arg):
   ...:         print(arg)
   ...: 
   ...:     @say.register
   ...:     def _(self, arg: int):
   ...:         print("Hi int", arg)
   ...: 
   ...:     @say.register
   ...:     def _(self, arg: list):
   ...:         print("Hi list", arg)
   ...: 

In [5]: h = Hi()

In [7]: h.say(1)
Hi int 1

In [8]: h.say("hello")
hello

In [9]: h.say([1, 2 ,3])
Hi list [1, 2, 3]

and similarly we can overload class methods:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
In [29]: class Hi:
    ...:     @singledispatchmethod
    ...:     @classmethod
    ...:     def say(cls, arg):
    ...:         print(arg)
    ...: 
    ...:     @say.register(int)
    ...:     @classmethod
    ...:     def _(cls, arg):
    ...:         print("Hi int", arg)
    ...: 
    ...:     @say.register(list)
    ...:     @classmethod
    ...:     def _(cls, arg):
    ...:         print("Hi list", arg)
    ...: 

In [30]: Hi.say(1)
Hi int 1

In [31]: Hi.say("hello")
hello

In [32]: Hi.say([1, 2,])
Hi list [1, 2]

In [33]: Hi.say([1, 2, 3])
Hi list [1, 2, 3]

Note that in order for the overload to work, singledispatchmethod and register must wrap classmethod. And that concludes today’s post!

Conclusion

Today we saw how to use singledispatch and singledispatchmethod from Python functools to overload functions and methods based on the first argument. This can come handy to refactor logic that would otherwise sit in the same implementation with if/else checking instance type, instead with singledispatch we can create multiple implementation and register them and the runtime will decide which implementation to call based on the first argument. I hope you liked this post and I’ll see you on the next one!

Designed, built and maintained by Kimserey Lam.