ProbableOdyssey

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()

Reply to this post by email ↪