TIL How to compose functions and coroutines in Python
A list of functions can be composed using this pattern:
1
2# Some example functions
3def foo(x):
4 return 3*x
5
6def bar(x):
7 return x + 3
8
9def baz(x):
10 return x**3
11
12# Function to compose functions
13def compose(*funcs):
14 def composition(x):
15 for func in funcs:
16 x = func(x)
17 return x
18 return composition
19
20# Example usage
21def main():
22 pipeline = compose(
23 foo,
24 bar,
25 baz
26 )
27
28 x = 5
29 x_out = pipeline(x)
30 print(x_out) # 5832
31
32if __name__ == "__main__":
33 main()
But what if one of your functions is a coroutine? And another one as a class
with an asynchronous __call__
method?
1import asyncio
2
3async def foo_async(x):
4 await asyncio.sleep(1)
5 return 3*x
6
7class BarAsync:
8 def __init__(self, floob):
9 self.floob = floob
10
11 async def __call__(self, x):
12 await asyncio.sleep(self.floob)
13 return 3*x
14
15def baz(x):
16 return x**3
Well we can modify our pipeline to handle these cases:
1def compose_async(*funcs):
2 async def composition_async(x):
3 for func in funcs:
4 result = func(x)
5 if asyncio.iscoroutine(result):
6 x = await result
7 else:
8 x = result
9 return x
10 return composition_async
11
12def main():
13 pipeline = compose_async(
14 foo_async,
15 BarAsync(2),
16 baz
17 )
18
19 x = 5
20 x_out = asyncio.run(pipeline(x)) # or `await` this in another coroutine
21 print(x_out) # 5832
22
23if __name__ == "__main__":
24 main()
Bonus tip: here’s how to type-hint the example to keep mypy
happy:
1import asyncio
2from typing import Union, Callable, Awaitable, cast
3
4# Pipeline functions
5async def foo_async(x: int) -> int:
6 await asyncio.sleep(1)
7 return 3*x
8
9class BarAsync:
10 def __init__(self, floob: int):
11 self.floob = floob
12
13 async def __call__(self, x: int) -> int:
14 await asyncio.sleep(self.floob)
15 return 3*x
16
17def baz(x: int) -> int:
18 return x**3
19
20
21# Define a type alias for functions in the pipeline
22PipelineFunc = Union[
23 Callable[[int], int],
24 Callable[[int], Awaitable[int]]
25]
26
27# Define pipeline
28def compose_async(*funcs: PipelineFunc) -> Callable[[int], Awaitable[int]]:
29 async def composition_async(x: int) -> int:
30 for func in funcs:
31 result = func(x)
32 if asyncio.iscoroutine(result):
33 x = await result
34 else:
35 x = cast(int, result) # Tell mypy explicitly about type
36 return x
37 return composition_async
38
39
40# Example of pipeline usage
41def main() -> None:
42 pipeline = compose_async(
43 foo_async,
44 BarAsync(2),
45 baz
46 )
47
48 x = 5
49 x_out = asyncio.run(pipeline(x)) # or `await` this in another coroutine
50 print(x_out) # 5832
51
52if __name__ == "__main__":
53 main()