Function Overload In Python With Single Dispatch
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!